module AlgebraHelpersTest

using Test
using ReactiveMP
using Random

using LinearAlgebra

@testset "Helpers" begin

    @testset "negate_inplace!" begin
        rng = MersenneTwister(1234)

        for size in 2:4
            A = rand(rng, size, size)
            B = copy(A)
            @test B == A
            @test ReactiveMP.negate_inplace!(B) == -A
            @test B == -A

            a = rand(rng)
            b = copy(a)
            @test ReactiveMP.negate_inplace!(b) == -a
            @test b == a # numbers stay unchanged
        end
    end

    @testset "negate_inplace!" begin
        rng = MersenneTwister(1234)

        for size in 2:4
            A = rand(rng, size, size)
            B = copy(A)
            c = randn(rng)
            @test B == A
            @test ReactiveMP.mul_inplace!(c, B) == c * A
            @test B == c * A

            a = rand(rng)
            b = copy(a)
            c = randn(rng)
            @test ReactiveMP.mul_inplace!(c, b) == c * a
            @test b == a # numbers stay unchanged
        end
    end

    @testset "rank1update" begin
        rng = MersenneTwister(1234)

        for size in 2:4, T1 in (Float32, Float64), T2 in (Float32, Float64), T3 in (Float32, Float64)
            A = rand(rng, T1, size, size)
            x = rand(rng, T2, size)
            @test ReactiveMP.rank1update(A, x) == (A + x*x')
            y = rand(rng, T3, size)
            @test ReactiveMP.rank1update(A, x, y) == (A + x*y')
        end

    end

    @testset "mul_trace" begin
        rng = MersenneTwister(1234)

        for size in 2:4, T1 in (Float32, Float64), T2 in (Float32, Float64)
            A = rand(rng, T1, size, size)
            B = rand(rng, T2, size, size)
            @test ReactiveMP.mul_trace(A, B) ≈ tr(A * B)

            a = rand(rng, T1)
            b = rand(rng, T2)
            @test ReactiveMP.mul_trace(a, b) ≈ a * b
        end

    end

end

end


include(joinpath(@__DIR__, "../../src/sketches/FindModel.jl"))
using Test

"""Define example FLS"""

#########################
# COMMUTATIVE TRIANGLE  #
#########################

f = FLSinit(@acset LabeledGraph begin
    V = 3
    E = 5
    vlabel = [:A,:B,:C]
    elabel = [:f,:h,:g, :ida, :inv]
    src    = [1,  1, 2,  1,    3]
    tgt    = [2,  3, 3,  1,    3]
end)

# f;g = h
add_diagram!(f, @acset LabeledGraph begin
    V = 3
    E = 3
    vlabel = [:A,:B,:C]
    elabel = [:f,:g,:h]
    src    = [1, 2, 1]
    tgt    = [2, 3, 3]
end)

# "ida" is an identity arrow
add_diagram!(f, @acset LabeledGraph begin
    V = 1
    E = 1
    vlabel = [:A]
    elabel = [:ida]
    src = [1]
    tgt = [1]
end);

# "inv" is an involution on C
add_diagram!(f, @acset LabeledGraph begin
    V = 2
    E = 2
    vlabel = [:C, :C]
    elabel = [:inv, :inv]
    src = [1, 2]
    tgt = [2, 1]
end)

#############################
# LEFT INVERSE / INVOLUTION #
#############################

leftinv = FLSinit(@acset LabeledGraph begin
    V = 2
    E = 4
    vlabel = [:A,:B]
    elabel = [:f,:g, :ida, :inv]
    src    = [1,  2,  1,   2]
    tgt    = [2,  1,  1,   2]
end);

add_diagram!(leftinv, @acset LabeledGraph begin
    V = 2
    E = 2
    vlabel = [:A, :B]
    elabel = [:f, :g]
    src = [1, 2]
    tgt = [2, 1]
end);

add_diagram!(leftinv, @acset LabeledGraph begin
    V = 1
    E = 1
    vlabel = [:A]
    elabel = [:ida]
    src = [1]
    tgt = [1]
end)

add_diagram!(leftinv, @acset LabeledGraph begin
    V = 2
    E = 2
    vlabel = [:B, :B]
    elabel = [:inv, :inv]
    src = [1, 2]
    tgt = [2, 1]
end)

#############
# EQUALIZER #
#############

fxb = Symbol("f×ᵦg")

equalizer = FLSinit(@acset LabeledGraph begin
    V = 3
    E = 4
    vlabel = [:A,:B,fxb]
    elabel = [:f,:g, :π1, :π2]
    src = [1,1,3,3]
    tgt = [2,2,1,1]
end)

add_cone!(equalizer, @acset LabeledGraph begin
    V = 4
    E = 4
    vlabel = [:A,:B,:A,fxb]
    elabel = [:f,:g, :π1, :π2]
    src = [1,3,4,4]
    tgt = [2,2,1,3]
end)

#################################
# SET PERMUTATIONS / PARTITIONS #
#################################

perm = FLSinit(@acset LabeledGraph begin
    V = 1
    E = 2
    vlabel = [:X]
    elabel = [:f, :f⁻¹]
    src = [1,1]
    tgt = [1,1]
end );

add_diagram!(perm, @acset LabeledGraph begin
    V = 2
    E = 2
    vlabel = [:X, :X]
    elabel = [:f, :f⁻¹]
    src = [1, 2]
    tgt = [2, 1]
end)

####################
# REFLEXIVE GRAPHS #
####################

reflg = FLSinit(@acset LabeledGraph begin
    V = 2
    E = 4
    vlabel = [:V, :E]
    elabel = [:id, :refl, :src, :tgt]
    src = [1,1,2,2]
    tgt = [1,2,1,1]
end );

add_diagram!(reflg, @acset LabeledGraph begin
    V = 2
    E = 4
    vlabel = [:V, :E]
    elabel = [:id, :refl, :src, :tgt]
    src = [1,1,2,2]
    tgt = [1,2,1,1]
end)

###########################
# CONSTANT (arrow from 1) #
###########################
s1 = Symbol("1")
const2 = FLSinit(@acset LabeledGraph begin
    V = 2
    E = 2
    vlabel = [s1, :A]
    elabel = [:f, :g]
    src = [1,1]
    tgt = [2,2]
end);

add_cone!(const2, @acset LabeledGraph begin
    V = 1
    vlabel=[s1]
end)

####################################
# Infinite sequence of a's and b's #
####################################

inflist = FLSinit(@acset LabeledGraph begin
    V = 3
    E = 4
    vlabel = [s1, :d, :l]
    elabel = [:a,:b,:head,:tail]
    src = [1,1,3,3]
    tgt = [2,2,2,3]
end )

add_cone!(inflist, @acset LabeledGraph begin
    V = 1
    vlabel = [s1]
end)

add_cone!(inflist, @acset LabeledGraph begin
    V = 3
    E = 2
    vlabel = [:d, :l, :l]
    elabel = [:head, :tail]
    src = [3,3]
    tgt = [1,2]
end)

#####################
# Cartesian Product #
#####################

pairs = FLSinit(@acset LabeledGraph begin
    V = 2
    E = 2
    vlabel = [:s, :s2]
    elabel = [:p1, :p2]
    src = [2,2]
    tgt = [1,1]
end ); # s s²

add_cone!(pairs, @acset LabeledGraph begin
    V = 3
    E = 2
    vlabel = [:s, :s, :s2]
    elabel = [:p1, :p2]
    src = [3,3]
    tgt = [1,2]
end )

srch_pairs = cs -> [catelems_to_cset(pairs, r) for r in search(pairs, cs)]

trips = FLSinit(@acset LabeledGraph begin
    V = 2
    E = 3
    vlabel = [:s, :s3]
    elabel = [:p1, :p2, :p3]
    src = [2,2,2]
    tgt = [1,1,1]
end ); # s s²

add_cone!(trips, @acset LabeledGraph begin
    V = 4
    E = 3
    vlabel = [:s, :s, :s, :s3]
    elabel = [:p1, :p2, :p3]
    src = [4,4,4]
    tgt = [1,2,3]
end )

##############
# SEMIGROUPS #
##############

p1p2, p2p3, idk, kid = map(Symbol, ["π₁×π₂","π₂×π₃","id×k","k×id"])
semig = FLSinit(@acset LabeledGraph begin
    V = 3
    E = 10
    vlabel = [:s, :s2, :s3]
    elabel = [:k, :π₁, :π₂, :Π₁, :Π₂, :Π₃, p1p2, p2p3, idk, kid]
    src    = [2,  2,   2,   3,    3,   3,   3,   3,    3,   3]
    tgt    = [1,  1,   1,   1,    1,   1,   2,   2,    2,   2]
end );

# Associativity
add_diagram!(semig, @acset LabeledGraph begin
    V = 4
    E = 4
    vlabel = [:s3, :s2, :s2, :s]
    elabel = [kid, idk, :k, :k]
    src = [1,1,2,3]
    tgt = [2,3,4,4]
end)

# Define p1p2
add_diagram!(semig, @acset LabeledGraph begin
    V = 4
    E = 5
    vlabel = [:s3, :s, :s, :s2]
    elabel = [:Π₁, :Π₂, p1p2, :π₁, :π₂]
    src = [1,1,1,4,4]
    tgt = [2,3,4,2,3]
end)

# Define p2p3
add_diagram!(semig, @acset LabeledGraph begin
    V = 4
    E = 5
    vlabel = [:s3, :s, :s, :s2]
    elabel = [:Π₂, :Π₃, p2p3, :π₁, :π₂]
    src    = [1,    1,   1,   4,   4]
    tgt    = [2,    3,   4,   2,   3]
end)

# Define kid
add_diagram!(semig, @acset LabeledGraph begin
    V = 5
    E = 6
    vlabel = [:s3, :s2,  :s2, :s, :s]

    elabel = [p1p2, kid, :Π₃, :π₂, :k, :π₁]
    src    = [1,    1,    1,   3,   2,  3]
    tgt    = [2,    3,    4,   4,   5,  5]
end)

# Define idk
add_diagram!(semig, @acset LabeledGraph begin
    V = 5
    E = 6
    vlabel = [:s3, :s2,  :s2, :s, :s]

    elabel = [p2p3, idk, :Π₁, :π₁, :k, :π₂]
    src    = [1,    1,    1,   3,   2,  3]
    tgt    = [2,    3,    4,   4,   5,  5]
end)

# s2 is pair
add_cone!(semig, @acset LabeledGraph begin
    V = 3
    E = 2
    vlabel = [:s, :s, :s2]
    elabel = [:π₁, :π₂]
    src = [3,3]
    tgt = [1,2]
end )

# s3 is triple
add_cone!(semig, @acset LabeledGraph begin
    V = 4
    E = 3
    vlabel = [:s, :s, :s, :s3]
    elabel = [:Π₁, :Π₂, :Π₃]
    src = [4,4,4]
    tgt = [1,2,3]
end )

# Colimits
import DataStructures: OrderedDict
function fls_pushout(l::LabeledFLS, r::LabeledFLS; kw...)::Fls
    ns, es = [OrderedDict([k=>v for (k, v) in kw if k in l.labels[i]])
              for i in [1,2]]
    overlap = Fls()
    add_parts!(overlap, :V, length(ns))
    add_parts!(overlap, :E, length(es))
    println("kw $kw Ns $ns Es $es")
    L = CSetTransformation(overlap, l.fls; V=[vi(l, x) for x in keys(ns)],
                                           E=[ei(l, x) for x in keys(es)])
    R = CSetTransformation(overlap, r.fls; V=[vi(r, x) for x in values(ns)],
                                           E=[ei(r, x) for x in values(es)])
    i, _ = pushout(L, R)
    return codom(i)
end

po = fls_pushout(reflg, equalizer; V=:B, E=:A, s=:f, t=:g)


"""
    compute_reduced_clique_graph!(sep::Array{Set{Int64}, 1}, snd::Array{Set{Int64}, 1})

Compute the reduced clique graph (union of all clique trees) given an initial clique tree defined by its supernodes `snd` and separator `sep` sets.

We are using the algorithm described in **Michel Habib and Juraj Stacho - Polynomial-time algorithm for the leafage ofchordal graphs (2009)**, which
computes the reduced clique graph in the following way:
1. Sort all minimal separators by size
2. Initialise graph CG(R) with cliques as nodes and no edges
3. for largest unprocessed separator S and
    |  add an edge between any two cliques C1 and C2 if they both contain S and are in different connected   components   of CG(R) and store in `edges`.
    |  Compute an edge weight used for merge decision and store in `val`.
    |  Store the index of the separator which is the intersection C1 ∩ C2 in `iter`
   end
"""
function compute_reduced_clique_graph!(sep::Array{Set{Int64}, 1}, snd::Array{Set{Int64}, 1})
    # loop over separators by decreasing cardinality
    sort!(sep, by = x -> length(x), rev = true)

    # edges = Array{Tuple{Int64, Int64}, 1}()    # a list of edges in the reduced clique graph, higher clique index first
    rows = Int64[]
    cols = Int64[]
    # inter = Array{Int64, 1}()                  # the index of the separator which corresponds to the intersection of the two cliques

    for (k, separator) in enumerate(sep)
        # find cliques that contain the separator
        clique_ind = findall(x -> separator ⊆ x, snd)

        # we compute the separator graph (see Habib, Stacho - Reduced clique graphs of chordal graphs) to analyse connectivity
        # we represent the separator graph H by a hashtable
        H = separator_graph(clique_ind, separator, snd)
        # find the connected components of H
        components = find_components(H, clique_ind)
        # for each pair of cliques that contain the separator, add an edge to the reduced clique tree if they are in unconnected components
        for pair in subsets(clique_ind, Val{2}())
            if is_unconnected(pair, components)
                # push!(edges, (max(pair...), min(pair...))) #add edge
                push!(rows, max(pair...))
                push!(cols, min(pair...))
                # push!(inter, k) # store intersection
            end
        end

    end
    return rows, cols#, inter
end

"Check whether the `pair` of cliques are in different `components`."
function is_unconnected(pair::Tuple{Int64, Int64}, components::Array{Set{Int64}, 1})
    component_ind = findfirst(x -> pair[1] ∈ x, components)
    return pair[2] ∉ components[component_ind]
end

"Find the separator graph H given a separator and the relevant index-subset of cliques."
function separator_graph(clique_ind::Array{Int64,1}, separator::Set{Int64}, snd::Array{Set{Int64}, 1})

    # make the separator graph using a hash table
    # key: clique_ind --> edges to other clique indices
    H = Dict{Int64, Array{Int64, 1}}()

    for pair in subsets(clique_ind, Val{2}())
        ca = pair[1]
        cb = pair[2]
         # if intersect_dim(snd[ca], snd[cb]) > length(separator)
        if !inter_equal(snd[ca], snd[cb], separator)
            if haskey(H, ca)
                push!(H[ca], cb)
            else
                H[ca] = [cb]
            end
            if haskey(H, cb)
                push!(H[cb], ca)
            else
                H[cb] = [ca]
            end
        end
    end
    # add unconnected cliques
    for v in clique_ind
        !haskey(H, v) && (H[v] = Int64[])
    end
    return H
end

"Find connected components in undirected separator graph represented by `H`."
function find_components(H::Dict{Int64, Array{Int64, 1}}, clique_ind::Array{Int64,1})
    visited = Dict{Int64, Bool}(v => false for v in clique_ind)
    components = Array{Set{Int64}, 1}()
    for v in clique_ind
        if visited[v] == false
            component = Set{Int64}()
            push!(components, DFS_hashtable!(component, v, visited, H))
        end
    end
    return components
end

"Depth first search on a hashtable `H`."
function DFS_hashtable!(component::Set{Int64}, v::Int64, visited::Dict{Int64, Bool}, H::Dict{Int64, Array{Int64, 1}})
    visited[v] = true
    push!(component, v)
    for n in H[v]
        if visited[n] == false
            DFS_hashtable!(component, n, visited, H)
        end
    end
    return component
end

"Check if s ∩ s2 == s3."
function inter_equal(s::Set{Int64}, s2::Set{Int64}, s3::Set{Int64})
    dim = 0
    len_s3 = length(s3)
    if length(s) < length(s2)
        sa = s
        sb = s2
    else
        sa = s2
        sb = s
    end
    for e in sa
        if e in sb
            dim += 1
            dim > len_s3 && return false
            e in s3 || return false
        end
    end
    return true
end

"Given a list of edges, return an adjacency hash-table `table` with nodes from 1 to `num_vertices`."
function compute_adjacency_table(edges::SparseMatrixCSC{Float64, Int64}, num_vertices::Int64)
    table = Dict(i => Set{Int64}() for i = 1:num_vertices)
    r = edges.rowval
    c = edges.colptr
     for col = 1:num_vertices
         for i in c[col]:c[col+1]-1
             row = r[i]
             push!(table[row], col)
            push!(table[col], row)
        end
     end
     return table
end

"Check whether `edge` is permissible for a merge. An edge is permissible if for every common neighbor N, C_1 ∩ N == C_2 ∩ N or if no common neighbors exist."
function ispermissible(edge::Tuple{Int64, Int64}, adjacency_table::Dict{Int64, Set{Int64}}, snd::Array{Set{Int64}, 1})
    c_1 = edge[1]
    c_2 = edge[2]
    common_neighbors = intersect(adjacency_table[c_1], adjacency_table[c_2])
    # N.B. This can be made faster by first checking whether the sizes of the intersection are the same before allocating anything
    for neighbor in common_neighbors
        intersect(snd[c_1], snd[neighbor]) != intersect(snd[c_2], snd[neighbor]) && return false
    end
    return true
end


@doc raw"""
`PSDE`: Stratonovich Partitioned Stochastic Differential Equation

Defines a partitioned stochastic differential initial value problem
```math
\begin{align*}
\dq (t) &= v(t, q(t)) \, dt + B(t, q(t)) \circ dW , & q(t_{0}) &= q_{0} ,
\dp (t) &= f(t, q(t)) \, dt + G(t, q(t)) \circ dW , & p(t_{0}) &= p_{0}
\end{align*}
```
with the drift vector fields ``v`` and ``f``, diffusion matrices ``B`` and ``G``,
initial conditions ``q_{0}`` and ``p_{0}``, the dynamical variables ``(q,p)`` taking
values in ``\mathbb{R}^{d} \times \mathbb{R}^{d}``, and the m-dimensional Wiener process W

### Fields

* `d`:  dimension of dynamical variable ``q`` and the vector field ``v``
* `m`:  dimension of the Wiener process
* `ni`: number of initial conditions
* `ns`: number of sample paths
* `v`:  function computing the drift vector field for the position variable q
* `f`:  function computing the drift vector field for the momentum variable p
* `B`:  function computing the d x m diffusion matrix for the position variable q
* `G`:  function computing the d x m diffusion matrix for the momentum variable p
* `t₀`: initial time
* `q₀`: initial condition for dynamical variable ``q`` (may be a random variable itself)
* `p₀`: initial condition for dynamical variable ``p`` (may be a random variable itself)

The functions `v`, `f`, 'B' and `G`, providing the drift vector fields and diffusion matrices, take four arguments,
`v(t, q, p, v)`, `f(t, q, p, f)`, `B(t, q, p,  B)` and `G(t, q, p, G)`, where `t` is the current time, `(q, p)` is the
current solution vector, and `v`, `f`, 'B' and `G` are the variables which hold the result
of evaluating the vector fields ``v``, ``f`` and the matrices ``B``, ``G`` on `t` and `(q,p)`.

### Example

```julia
    function v(λ, t, q, v)
        v[1] = λ*q[1]
        v[2] = λ*q[2]
    end

    function B(μ, t, q, B)
        B[1] = μ*q[1]
        B[2] = μ*q[2]
    end

    t₀ = 0.
    q₀ = [1., 1.]
    λ  = 2.
    μ  = 1.

    v_sde = (t, q, v) -> v(λ, t, q, v)
    B_sde = (t, q, B) -> B(μ, t, q, B)

    sde = SDE(v_sde, B_sde, t₀, q₀)
```
"""
struct PSDE{dType <: Number, tType <: Number, vType <: Function, fType <: Function,
            BType <: Function, GType <: Function, pType <: Union{Tuple,Nothing}, N} <: AbstractEquationPSDE{dType, tType}

    d::Int
    m::Int
    ni::Int
    ns::Int
    v::vType
    f::fType
    B::BType
    G::GType
    t₀::tType
    q₀::Array{dType, N}           #Initial condition: N=1 - single deterministic, N=2 - single random or multiple deterministic, N=3 - multiple deterministic
    p₀::Array{dType, N}
    parameters::pType
    periodicity::Vector{dType}

    function PSDE(m, ns, v::vType, f::fType, B::BType, G::GType,
                  t₀::tType, q₀::DenseArray{dType}, p₀::DenseArray{dType};
                  parameters=nothing, periodicity=zeros(dType,size(q₀,1))) where {
                        dType <: Number, tType <: Number,
                        vType <: Function, fType <: Function,
                        BType <: Function, GType <: Function}

        @assert size(q₀)  == size(p₀)
        @assert ndims(q₀) == ndims(p₀)

        N  = ndims(q₀)
        d  = size(q₀,1)
        ni = size(q₀,2)

        @assert N ∈ (1,2)
        @assert ni ≥ 1
        @assert ns ≥ 1
        @assert ni == 1 || ns == 1
        # either multiple deterministic initial conditions and one sample path
        # or one deterministic initial condition and multiple sample paths

        new{dType,tType,vType,fType,BType,GType,typeof(parameters),N}(d, m, ni, ns, v, f, B, G, t₀, q₀, p₀, parameters, periodicity)
    end
end

function PSDE(m::Int, ns::Int, v::Function, f::Function, B::Function, G::Function, q₀::DenseArray{DT}, p₀::DenseArray{DT}; kwargs...) where {DT}
    PSDE(m, ns, v, f, B, G, zero(DT), q₀, p₀; kwargs...)
end

Base.hash(sde::PSDE, h::UInt) = hash(sde.d, hash(sde.m, hash(sde.ni, hash(sde.ns,
                                hash(sde.v, hash(sde.f, hash(sde.B, hash(sde.G,
                                hash(sde.t₀, hash(sde.q₀, hash(sde.p₀,
                                hash(sde.parameters, hash(sde.periodicity, h)))))))))))))

Base.:(==)(sde1::PSDE, sde2::PSDE) = (
                                sde1.d == sde2.d
                             && sde1.m == sde2.m
                             && sde1.ni == sde2.ni
                             && sde1.ns == sde2.ns
                             && sde1.v == sde2.v
                             && sde1.f == sde2.f
                             && sde1.B == sde2.B
                             && sde1.G == sde2.G
                             && sde1.t₀ == sde2.t₀
                             && sde1.q₀ == sde2.q₀
                             && sde1.p₀ == sde2.p₀
                             && sde1.parameters  == sde2.parameters
                             && sde1.periodicity == sde2.periodicity)

function Base.similar(sde::PSDE, q₀::AbstractArray, p₀::AbstractArray, ns::Int)
    similar(sde, sde.t₀, q₀, p₀, ns)
end

function Base.similar(sde::PSDE, q₀::AbstractArray, p₀::AbstractArray)
    similar(sde, sde.t₀, q₀, p₀)
end

function Base.similar(sde::PSDE, t₀::TT, q₀::AbstractArray{DT,N}, p₀::AbstractArray{DT,N}, ns::Int=(N > 1 ? 1 : sde.ns)) where {DT <: Number, TT <: Number, N}
    @assert size(q₀) == size(p₀)
    @assert sde.d == size(q₀,1) == size(p₀,1)
    PSDE(sde.m, ns, sde.v, sde.f, sde.B, sde.G, t₀, q₀, p₀; parameters=sde.parameters, periodicity=sde.periodicity)
end

@inline Base.ndims(sde::PSDE) = sde.d

@inline periodicity(equation::PSDE) = equation.periodicity


using Roots
using Statistics
using FileIO
using Parameters
include("random_matrices.jl")
include("network_sims_gpu.jl")

etas = collect(-1.0f0:0.1f0:1.0f0)
ind = parse(Int64, ARGS[1])

N = 3000
eta = etas[ind]
s = 0.0f0

J_seeds = collect(5000:5019)
h0_seed = 6000

h0 = generate_nv(N, Float32, seed=h0_seed)

wait = 200.0f0
t_bound = 2000.0f0
dt = 0.01f0

startstep = round(Int, wait / dt) + 1
lateststep = round(Int, t_bound / dt)

@unpack target_avg_var = load("vanilla_target_avg_var.jld2")

# Difference between the average variance of the network ensemble and target
function f(g)
    vanilla_network_avg_vars = Vector{Float64}(undef, length(J_seeds))
    for (i, J_seed) in enumerate(J_seeds)
        J = generate_J(N, eta, seed=J_seed)
        A = sim_activity_gpu(J, g, h0; s=s, dt=dt, t_bound=t_bound,
                minimum_norm=1e-90)
        endstep = min(lateststep, size(A, 2))
        avg_var = mean(var(view(A, :, startstep:endstep), dims=2))
        vanilla_network_avg_vars[i] = avg_var
    end
    mean(vanilla_network_avg_vars) - target_avg_var
end

lo_bound = 0.0f0
hi_bound = 100f0
@time g = find_root(f, (lo_bound, hi_bound), Bisection(), verbose=true)

out = Dict{String,Any}()
@pack! out = N, g, eta, s, J_seeds, h0_seed, wait, t_bound, dt, target_avg_var
save("Rootfinding_E$(ind)_N=$N,eta=$eta,s=$s.jld2", out)


function search_frequency(rest_ν, ν)
  min_i = 1
  min_Δ = Inf
  for (i, value) in enumerate(ν)
    current_Δ = abs(value - rest_ν)
    (min_Δ, min_i) = current_Δ < min_Δ ? (current_Δ, i) : (min_Δ, min_i)
  end
  return min_Δ, min_i
end

"""Loops through two vectors A and B, searching for the best matches
for between them.

Matches are only considered if the difference of the values meet
a preliminary threshold, as a fraction of the value of A.

This is the preferred method of matching the catalog and obs. frequencies,
as it means you only run the loop once.

Returns a list of index pairs of A and B.
"""
function intersect_frequencies(A::Vector{T}, B::Vector{S}; thres=1f-4) where {T<:Real,S<:Real}
  matches = []
  # preallocate up to the length of the rest frequencies
  sizehint!(matches, length(A))
  for (i, rest_value) in enumerate(A)
    to_match = rest_value * thres
    best_diff = Inf
    best_match = nothing
    for (j, obs_value) in enumerate(B)
      diff = abs(rest_value - obs_value)
      if diff <= to_match && diff < best_diff
        best_diff = diff
        best_match = [i, j]
      end
      # if the diff increases from the best diff, we've passed it
      if diff > best_diff
        push!(matches, best_match)
        break
      end
    end
  end
  return matches
end

#function extract_windows(rest_ν, obs::Observation, vlsr; thres=1f-4,)
#  matches = intersect_frequencies(rest_ν, obs.frequency; thres=thres)
#end


"""
  outer!(z::Vector{Float}, a::Vector{Float}, b::Vector{Float})

Return z - Vector{Float} , vectorsed outer/kroneker product o vectors a and b
Auxiliary function for rawmoment
"""
function outer!(z::Vector{T}, a::Vector{T}, b::Vector{T}) where T<: AbstractFloat
    sa = size(a,1)
    sb = size(b,1)
    m = 1
    for i = 1:sb for j = 1:sa
        @inbounds z[m] = a[j]*b[i]
        m += 1
    end
    end
end

"""
  updvec!(A::Vector{Float}, B::Vector{Float})

Returns updated Vector{Float} A, by adding elementwisely Vector{Float} B
Auxiliary function for rawmoment
"""
function updvec!(A::Vector{T}, B::Vector{T}) where T<: AbstractFloat
  n = size(A, 1)
  for i=1:n
    @inbounds A[i] += B[i]
  end
  return A
end

"""
  rawmoment(X::Matrix{T}, m::Int = 4)

Simmilar to raw_moments_upto_p in R, does not expoloit tensor's symmetry
pyramid structures and blocks

Returns Array{Float, m}, the m'th moment's tensor
"""
function rawmoment(X::Matrix{T}, m::Int = 4) where T<: AbstractFloat
  t,n = size(X)
  if m == 1
    return mean(X, dims=1)[1,:]
  else
    z = [map(i -> zeros(T, n^i), 1:m)...]
    y = zeros(T, n^m)
    for i in 1:t
      xi = X[i, :]
      z[1] = xi
      for j in 2:m
        outer!(z[j], xi, z[j-1])
      end
      updvec!(y, z[m])
    end
  end
  reshape(y/t, fill(n, m)...)
end

"""
  raw_moments_upto_k(X::Matrix, k::Int = 4)

Returns [Array{Float, 1}, ..., Array{Float, k}] noncentral moment tensors of
order 1, ..., k
"""
raw_moments_upto_k(X::Matrix{T}, k::Int = 4) where T<: AbstractFloat =
  [rawmoment(X, i) for i in 1:k]

"""
  cumulants_from_moments(raw::Vector{Array{Float, i}, m = 1:k})

Returns [Array{Float, 1}, ..., Array{Float, k}] cumulant tensors of order 1, ..., k

Uses relation between cumulants and multivariate moments from e.g.
c_{ijkl} = m_{ijkl} - m_{ijk} m_{l} [4] - m_{ij} m_{kl} [3] + 2 m_{ij} m_{k} m_{l}
-6 m_{i} m_{j} m_{k} m_{l}
"""
function cumulants_from_moments(raw::Vector{Array{T}}) where T<: AbstractFloat
  k = length(raw)
  cumarr = Array{Array{Float64}}(undef, k)
  for j in 1:k
    dimr = size(raw[j])
    cumu = zeros(Float64, dimr)
    ldim = length(dimr)
    spp = collect(partitions(1:ldim))
    qpp = [(-1)^i*factorial(i) for i in 0:(ldim-1)]
    sppl = [map(length, spp[i]) for i in 1:length(spp)]
    for i in 1:prod(dimr)
      @inbounds ind = Tuple(CartesianIndices(dimr)[i])
      @inbounds cumu[ind...] = onecumulant(ind, raw, spp, sppl, qpp)
    end
    cumarr[j] = cumu
  end
  cumarr
end

"""
  onecumulant(ind::Tuple, raw::Vector{Array}, spp::Vector, sppl::Vector{Vector}, dpp::Vector)

raw - vector of moment's tensors, spp - vector of partitions, sppl - vector of sizes of partitions
dpp - vector of a factor for each product of moments (a factorial factor).

Returns Array{Float, n} the n'th cumulant tensor

"""
function onecumulant(ind::Tuple, raw::Vector{Array{T}}, spp::Vector,
          sppl::Vector{Vector{Int}}, dpp::Vector{Int}) where T<: AbstractFloat
  ret = zero(T)
  for i in 1:length(spp)
    part = spp[i]
    beln = length(part)
    k = sppl[i]
    temp = one(T)
    for r in 1:beln
      temp *= raw[k[r]][ind[part[r]]...]
    end
    ret += dpp[beln]*temp
  end
  ret
end

"""
  cumulatsfrommoments(x::Matrix{Float}, k::Int)

Returns a vector of 1,2, .., k dims Arrays{Float} of cumulant tensors
"""
mom2cums(x::Matrix{T}, k::Int) where T<: AbstractFloat =
  cumulants_from_moments(raw_moments_upto_k(x, k))


########################### Chisq Diagnostic ###########################

function update_hangartner_temp_vars!(u::Matrix{Int64}, X::Matrix{Int64},
                                      t::Int64)
  d = size(X,2)

  for j in 1:d
    u[X[t, j], j] += 1
  end
end

function hangartner_inner(Y::AbstractMatrix, m::Int)
  ## setup temp vars
  n, d = size(Y)

  # Count for each category in each chain
  u = zeros(Int64, m, d)
  v = zeros(Int64, m, d)

  for t in 1:n
    # fill out temp vars
    update_hangartner_temp_vars!(u, Y, t)
  end
  phia, chi_stat, m_tot = weiss_sub(u, v, n)

  return (n * sum(chi_stat), m_tot)
end

function weiss(X::AbstractMatrix{U}) where {U<:Any}
  ## number of iterations, number of chains
  n, d = size(X)

  ## mapping of values to integers
  v_dict = Dict{U, Int}()

  ## max unique categories
  mc = map(c -> length(unique(X[:,c])), 1:d)
  m = length(unique(X))

  ## counter of number of unique values in each chain
  r0 = 0

  ## Count for each category in each chain
  u = zeros(Int, m, d)

  ## Number of times a category did not transition in each chain
  v = zeros(Int, m, d)

  for t in 1:n
    for c in 1:d
      if !(X[t,c] in keys(v_dict))
        r0 += 1
        v_dict[X[t,c]] = r0
      end
      idx1 = v_dict[X[t,c]]
      u[idx1,c] += 1

      if t > 1
        if X[t-1,c] == X[t,c]
          v[idx1, c] += 1
        end
      end
    end
  end
  phia, chi_stat, m_tot = weiss_sub(u, v, n)
  ca = (1 + phia) / (1 - phia)
  stat = (n / ca) * sum(chi_stat)
  pval = NaN
  if ((m_tot - 1) * (d - 1)) >= 1
    pval = 1 - cdf(Chisq((m_tot - 1) * (d - 1)), stat)
  end

  return ( stat, m_tot, pval, ca)
end

function weiss_sub(u::Matrix{Int64}, v::Matrix{Int64}, t::Int)
  m, d = size(u)
  nt = 0.0
  dt = 0.0
  m_tot = 0

  mp = zeros(Float64, m, d)
  ma = zeros(Float64, m)
  phia = 0.0
  ca = 0.0
  df = 0.0

  chi_stat = zeros(Float64, d)

  for j in 1:m
    p1 = 0.0
    p2 = 0.0
    for l in 1:d
      #aggregate
      p1 += v[j, l] / (d * (t - 1))
      p2 += u[j, l] / (d * t)

      #per chain
      mp[j, l] = u[j, l] / t
      ma[j] += u[j, l] / (d * t)
    end
    nt += p1
    dt += p2 ^ 2

    if ma[j] > 0
      m_tot += 1
      for l in 1:d
        chi_stat[l] += (mp[j,l] - ma[j]) ^ 2 / ma[j]
      end
    end
  end
  phia = 1.0 + (1.0 / t) - ((1 - nt) / (1 - dt))
  phia = min(max(phia,0.0),1.0 - eps())
  return (phia, chi_stat, m_tot)
end

function update_billingsley_temp_vars!(f::Array{Int64,3},
                                       X::Matrix{Int64},
                                       t::Int64)
  d = size(X,2)
  for j in 1:d
    if t > 1
      f[X[t-1, j], X[t,j], j] += 1
    end
  end
end

function billingsley_sub(f::Array{Int64,3})
  df = 0.0
  stat = 0.0

  m, d = size(f, 2), size(f, 3)

  # marginal transitions, i.e.
  # number of transitions from each category
  mf = mapslices(sum, f, dims=2)

  # For each category, number of chains for which
  # that category was present
  A = vec(mapslices( (x)-> sum(x.>0), mf, dims=3))

  # For each category, number of categories it
  # transitioned to
  B =  vec(mapslices((x)-> sum(x.>0), mapslices(sum,f,dims=3),dims=2))

  # transition probabilities in each chain
  P = f ./ mf

  # transition probabilities
  mP = reshape((mapslices(sum,f,dims=3) ./ mapslices(sum,mf,dims=3)), Val(2))

  idx = findall(x -> x != 0, A .* B)
  for j in idx
    #df for billingsley
    df += (A[j] - 1) * (B[j] - 1)

    #billingsley
    for k in idx
      if (mP[j,k] > 0.0)
        for l in 1:d
          if mf[j,1,l] > 0 && isfinite(P[j,k,l])
            stat += mf[j,1,l] * (P[j,k,l] - mP[j,k]) ^ 2 / mP[j,k]
          end
        end
      end
    end
  end
  return (stat, df, mP)
end

function bd_inner(Y::AbstractMatrix, m::Int)
  num_iters, num_chains = size(Y)
  # Transition matrix for each chain
  f = zeros(Int64, m, m, num_chains)

  for t in 1:num_iters
    # fill out temp vars
    update_billingsley_temp_vars!(f, Y, t)
  end
  billingsley_sub(f)
end

function simulate_NDARMA(N::Int, p::Int, q::Int, prob::Vector{Float64},
                         phi::Vector{Float64})
  X = zeros(Int64, N)
  X[1:p] = rand(Categorical(prob), p)
  d1 = Multinomial(1, phi)
  d2 = Categorical(prob)
  for t in (p+1):N
    alphabeta = rand(d1)
    eps = rand(d2, q + 1)
    X[t] = sum([X[(t-p):(t-1)]; eps] .* alphabeta)
  end
  return X
end

function simulate_MC(N::Int, P::Matrix{Float64})
  X = zeros(Int64, N)
  n, m = size(P)
  X[1] = sample(1:n)
  for i in 2:N
    X[i] = wsample(1:n, vec(P[X[i-1], :]) )
  end
  return X
end

function diag_all(X::AbstractMatrix{U}, method::Symbol,
                  nsim::Int, start_iter::Int, step_size::Int) where {U<:Any}

  ## number of iterations, number of chains
  n, d = size(X)

  ## mapping of values to integers
  v_dict = Dict{U, Int}()

  ## max unique categories
  mc = map(c -> length(unique(X[:,c])), 1:d)
  m = length(unique(X))

  ## counter of number of unique values in each chain
  r0 = 0

  ## Count for each category in each chain
  u = zeros(Int, m, d)

  ## Number of times a category did not transition in each chain
  v = zeros(Int, m, d)

  ## transition matrix for each chain
  f = zeros(Int, m, m, d)

  result = zeros(Float64, 3, length(start_iter:step_size:n))
  result_iter = 1
  for t in 1:n
    for c in 1:d
      if !(X[t,c] in keys(v_dict))
        r0 += 1
        v_dict[X[t,c]] = r0
      end
      idx1 = v_dict[X[t,c]]
      u[idx1,c] += 1

      if t > 1
        idx2 = v_dict[X[t-1,c]]
        f[idx1, idx2, c] += 1

        if X[t-1,c] == X[t,c]
          v[idx1, c] += 1
        end
      end
    end

    if ((t >= start_iter) && (((t - start_iter) % step_size) == 0))
      phia, chi_stat, m_tot = weiss_sub(u, v, t)
      hot_stat, df, mP = billingsley_sub(f)

      phat = mapslices(sum,u,dims=2)[:,1] / sum(mapslices(sum,u,dims=2))
      ca = (1 + phia) / (1 - phia)
      stat = NaN
      pval = NaN
      df0  = NaN

      if method == :hangartner
        stat = t * sum(chi_stat)
        df0 = (m - 1) * (d - 1)
        if m > 1 && !isnan(stat)
          pval = 1 - cdf(Chisq( (m - 1) * (d - 1)), stat)
        end
      elseif method == :weiss
        stat = (t / ca) * sum(chi_stat)
        df0 = (m - 1) * (d - 1)
        pval = NaN
        if m > 1 && !isnan(stat)
          pval = 1 - cdf(Chisq( (m - 1) * (d - 1)), stat)
        end
      elseif method == :DARBOOT
        stat = t * sum(chi_stat)
        bstats = zeros(Float64, nsim)
        for b in 1:nsim
          Y = hcat([simulate_NDARMA(t, 1, 0, phat, [phia, 1-phia])
                for j in 1:d]...)
          s = hangartner_inner(Y, m)[1]
          bstats[b] = s
        end
        idx = findall(!isnan(bstats))
        df0 = mean(bstats[idx])
        pval = sum(stat .<= bstats[idx])/length(idx)
      elseif method == :MCBOOT
        bstats = zeros(Float64, nsim)
        for b in 1:nsim
          Y = hcat([simulate_MC(t, mP) for j in 1:d]...)
          s = hangartner_inner(Y, m)[1]
          bstats[b] = s
        end
        idx = findall(!isnan(bstats))
        df0 = mean(bstats[idx])
        pval = sum(stat .<= bstats[idx])/length(idx)
      elseif method == :billingsley
        stat = hot_stat
        df0 = df
        if df > 0 && !isnan(hot_stat)
          pval = 1 - cdf(Chisq(df), hot_stat)
        end
      elseif method == :billingsleyBOOT
        stat = hot_stat
        bstats = zeros(Float64, nsim)
        for b in 1:nsim
          Y = hcat([simulate_MC(t, mP) for j in 1:d]...)
          (s,sd) = bd_inner(Y, m)[1:2]
          bstats[b] = s/sd
        end
        idx = findall(!isnan(bstats))
        df0 = mean(bstats[idx])
        pval = sum(stat/df .<= bstats[idx])/length(idx)
      else
        error("Unexpected")
      end
      result[:, result_iter] = [stat, df0, pval]
      result_iter += 1
    end
  end
  return result
end

function discretediag_sub(c::AbstractChains, frac::Real, method::Symbol,
                          nsim::Int, start_iter::Int, step_size::Int)

  num_iters, num_vars, num_chains = size(c.value)

  vals = zeros(Float64, 3 * (num_chains + 1), num_vars)
  plot_vals_stat = zeros(length(start_iter:step_size:num_iters), num_vars)
  plot_vals_pval = zeros(length(start_iter:step_size:num_iters), num_vars)

  ## Between-chain diagnostic
  X = zeros(Int64, num_iters, num_chains)
  for j in 1:length(num_vars)
    X = convert(Array{Int64, 2}, c.value[:,j,:])
    result = diag_all(X, method, nsim, start_iter, step_size)
    plot_vals_stat[:,j] = result[1, :] ./ result[2, :]
    plot_vals_pval[:,j] = result[3, :]
    vals[1:3, j] = result[:, end]
  end

  ## Within-chain diagnostic
  x = zeros(Int64, num_iters)
  Y = zeros(Int64, num_iters, 2)
  for j in 1:num_vars
    for k in 1:num_chains
      x = convert(Array{Int64, 1}, c.value[:,j,k])

      idx1 = 1:round(Int64, frac * num_iters)
      idx2 = round(Int64, num_iters - frac * num_iters + 1):num_iters
      x1 = x[idx1]
      x2 = x[idx2]
      n_min = min(length(x1), length(x2))
      Y = [x1[1:n_min] x2[(end - n_min + 1):end]]

      vals[(3 + 3 * (k - 1) + 1):(3 + 3 * (k - 1) + 3), j] =
        diag_all(Y, method, nsim, n_min, step_size)[:, end]
    end
  end
  return (collect(1:num_vars), vals, plot_vals_stat, plot_vals_pval)

end

function discretediagplot(c::AbstractChains; frac::Real=0.3,
                          method::Symbol=:weiss, nsim::Int=1000,
                          start_iter::Int=100, step_size::Int=10000)

  num_iters, num_vars, num_chains = size(c.value)

  valid_methods = [:hangartner, :weiss, :DARBOOT,
                   :MCBOOT, :billingsley, :billingsleyBOOT]
  if !(method in valid_methods)
    methods_str = join([":$f" for f in valid_methods], ", ")
    throw(ArgumentError("method must be one of ", methods_str))
  end

  if !(0.0 < frac < 1.0)
    throw(ArgumentError("frac must be in (0,1)"))
  end

  if (start_iter > num_iters ) || (step_size > num_iters)
    throw(ArgumentError("start_iter, step_size must be less than $num_iters"))
  end

  V, vals, plot_vals_stat, plot_vals_pval =
    discretediag_sub(c, frac, method, nsim, start_iter, step_size)

  p1 = plot(y=vcat([plot_vals_stat[:,j] for j in 1:length(V)]...),
            x=repeat(collect(c.range[start_iter:step_size:num_iters])/1000,
                     outer=[length(V)]),
            Geom.line,
            Guide.xlabel("Iteration (thousands)", orientation=:horizontal),
            Guide.ylabel("stat/df",orientation=:vertical),
            Scale.color_discrete(), Guide.colorkey(title="Variable"),
            color=repeat(c.names[V],
                         inner=[length(start_iter:step_size:num_iters)]))

  p2 = plot(y=vcat([plot_vals_pval[:,j] for j in 1:length(V)]...),
            x=repeat(collect(c.range[start_iter:step_size:num_iters])/1000,
                     outer=[length(V)]),
            Geom.line,
            Guide.xlabel("Iteration (thousands)", orientation=:horizontal),
            Guide.ylabel("pval",orientation=:vertical),
            Scale.color_discrete(), Guide.colorkey(title="Variable"),
            color=repeat(c.names[V],
                         inner=[length(start_iter:step_size:num_iters)]))

  return [p1, p2]
end

function discretediag(c::AbstractChains; frac::Real=0.3,
                      method::Symbol=:weiss, nsim::Int=1000)

  num_iters, num_vars, num_chains = size(c.value)

  valid_methods = [:hangartner, :weiss, :DARBOOT,
                   :MCBOOT, :billingsley, :billingsleyBOOT]
  if !(method in valid_methods)
    methods_str = join([":$f" for f in valid_methods], ", ")
    throw(ArgumentError("method must be one of ", methods_str))
  end

  if !(0.0 < frac < 1.0)
    throw(ArgumentError("frac must be in (0,1)"))
  end

  V, vals = discretediag_sub(c, frac, method, nsim,
                             size(c.value,1), size(c.value,1))[1:2]

  hdr = header(c) * "\nChisq Diagnostic:\nEnd Fractions = $frac\n" *
  "method = $method\n"
  ChainSummary(copy(round.(vals,digits=3)'), c.names[V],
               convert(Array{AbstractString, 1},
                       vcat([["stat", "df", "p-value"]
                             for k in 1:(num_chains + 1)]...)), hdr)
end


"""
Type: LSAtT

Methods: update!, reinit!, copy

A structure designed to track line search information from one iteration to
another. If we have f : ℜⁿ → ℜ, then we define h(θ) = f(x + θ*d) where x and d
are vectors of same dimension and θ is a scalar, more specifically our step size.

Tracked data can include:
 - x : our current step size
 - ht : h(θ) at the current iteration
 - gt : h'(θ)
 - h₀ : h(0)
 - g₀ : h'(0)
 - current_time:  the time at which the line search algorithm started.
 - current_score: the score at which the line search algorithm started.

Note: by default, unknown entries are set to nothing.
"""
mutable struct 	LSAtT <: AbstractState

    x            :: Number
    ht           :: FloatVoid  # h(θ)
    gt           :: FloatVoid  # h'(θ)
    h₀           :: FloatVoid  # h(0)
    g₀           :: FloatVoid  # h'(0)

    current_time   :: FloatVoid
    current_score  :: FloatVoid

 function LSAtT(t             :: Number;
                ht            :: FloatVoid = nothing,
                gt            :: FloatVoid = nothing,
                h₀            :: FloatVoid = nothing,
                g₀            :: FloatVoid = nothing,
                current_time  :: FloatVoid = nothing,
                current_score :: FloatVoid = nothing)

  return new(t, ht, gt, h₀, g₀, current_time, current_score)
 end
end

import Base.copy
"""
copy: Copy a LSAtT
"""
function copy(ls_at_t :: LSAtT)
    return LSAtT(copy(ls_at_t.x),
                 ht = copy(ls_at_t.ht),
                 gt = copy(ls_at_t.gt),
                 h₀ = copy(ls_at_t.h₀),
                 g₀ = copy(ls_at_t.g₀),
                 current_time = copy(ls_at_t.current_time),
                 current_score = copy(ls_at_t.current_score))
end


function solve(days, counts)
    for day in 1:days
        first = popfirst!(counts)
        push!(counts, 0)
        counts[7] += first
        counts[9] += first
    end
    return sum(counts)
end

nums = (x -> parse(Int, x)).(split(readline(), ','))
counts = [sum(nums .== i) for i in 0:9]

println(solve(80, copy(counts)))
println(solve(256, copy(counts)))


"""
    functional_cartography(N::T, L::Dict{E,Int64}) where {T<:BinaryNetwork, E<:AllowedSpeciesTypes}

This function will take the output of a modularity analysis (*i.e.* a network
and a partition), and return a dictionary where every species is associated to
its functional role, as defined in Olesen et al (2005). The first element is the
within-module degree z-score, and the second is the participation coefficient.
"""
function functional_cartography(N::T, L::Dict{E,Int64}) where {T<:BinaryNetwork, E<:AllowedSpeciesTypes}
    κ = Dict{last(eltype(N)), Int64}()
    S = Dict{Int64,Tuple{Float64, Float64}}()
    carto = Dict{last(eltype(N)),Tuple{Float64, Float64}}()

    for s in species(N)
        sp_module = L[s]
        k_i = 0
        if s ∈ species(N; dims=2)
            for s2 in N[:,s]
                L[s2] == sp_module && (k_i += 1)
            end
        end
        if s ∈ species(N; dims=1)
            for s2 in N[s,:]
                L[s2] == sp_module && (k_i += 1)
            end
        end
        κ[s] = k_i
    end

    for s in values(L)
        sp = filter(k -> k.second == s, L)
        k_i = [get(κ, x, 0) for x in keys(sp)]
        S[s] = (mean(k_i), std(k_i))
    end

    # Participation coefficient
    for s in species(N)
        ks = zeros(Int64, length(unique(values(L))))
        if s ∈ species(N; dims=2)
            for s2 in N[:,s]
                ks[L[s2]] += 1
            end
        end
        if s ∈ species(N; dims=1)
            for s2 in N[s,:]
                ks[L[s2]] += 1
            end
        end
        carto[s] = ((κ[s] - S[L[s]][1])/S[L[s]][2], 1-sum((ks./degree(N)[s]).^2))
    end

    return carto
end


using Test
using LinearAlgebra
using OrdinaryDiffEq, DiffEqDevTools

# Since non-FSAL methods are FSAL'ed artificially in OrdinaryDiffEq.jl,
# we need to exclude this last FSAL stage for the comparison with
# tableaus.
function coefficients_as_in_tableau(A, b, c, tab)
  if size(A) == size(tab.A)
    AA = A
    bb = b
    cc = c
  else
    AA = A[1:end-1, 1:end-1]
    bb = b[1:end-1]
    cc = c[1:end-1]
  end
  AA, bb, cc
end

function ode_tableau_tests(T)
  let erk = Heun()
    A, b, c = deduce_Butcher_tableau(erk)
    tab = constructHeun()
    AA, bb, cc = coefficients_as_in_tableau(A, b, c, tab)
    @test AA ≈ tab.A
    @test bb ≈ tab.α
    @test cc ≈ tab.c
  end

  let erk = Euler()
    A, b, c = deduce_Butcher_tableau(erk)
    tab = constructEuler()
    AA, bb, cc = coefficients_as_in_tableau(A, b, c, tab)
    @test AA ≈ tab.A
    @test bb ≈ tab.α
    @test cc ≈ tab.c
  end

  let erk = RK4()
    A, b, c = deduce_Butcher_tableau(erk)
    tab = constructRK4()
    AA, bb, cc = coefficients_as_in_tableau(A, b, c, tab)
    @test AA ≈ tab.A
    @test bb ≈ tab.α
    @test cc ≈ tab.c
  end

  let erk = SSPRK22()
    A, b, c = deduce_Butcher_tableau(erk)
    tab = constructSSPRK22()
    AA, bb, cc = coefficients_as_in_tableau(A, b, c, tab)
    @test AA ≈ tab.A
    @test bb ≈ tab.α
    @test cc ≈ tab.c
  end

  let erk = SSPRK33()
    A, b, c = deduce_Butcher_tableau(erk)
    tab = constructSSPRK33()
    AA, bb, cc = coefficients_as_in_tableau(A, b, c, tab)
    @test AA ≈ tab.A
    @test bb ≈ tab.α
    @test cc ≈ tab.c
  end

  let erk = SSPRK432()
    A, b, c = deduce_Butcher_tableau(erk)
    tab = constructSSPRK43()
    AA, bb, cc = coefficients_as_in_tableau(A, b, c, tab)
    @test AA ≈ tab.A
    @test bb ≈ tab.α
    @test cc ≈ tab.c
  end

  let erk = SSPRK104()
    A, b, c = deduce_Butcher_tableau(erk)
    tab = constructSSPRK104()
    AA, bb, cc = coefficients_as_in_tableau(A, b, c, tab)
    @test AA ≈ tab.A
    @test bb ≈ tab.α
    @test cc ≈ tab.c
  end

  let erk = Tsit5()
    A, b, c = deduce_Butcher_tableau(erk)
    tab = constructTsitouras5()
    AA, bb, cc = coefficients_as_in_tableau(A, b, c, tab)
    @test AA ≈ tab.A
    @test bb ≈ tab.α
    @test cc ≈ tab.c
  end

  let erk = BS3()
    A, b, c = deduce_Butcher_tableau(erk)
    tab = constructBogakiShampine3()
    AA, bb, cc = coefficients_as_in_tableau(A, b, c, tab)
    @test AA ≈ tab.A
    @test bb ≈ tab.α
    @test cc ≈ tab.c
  end

  let erk = BS5()
    A, b, c = deduce_Butcher_tableau(erk)
    tab = constructBogakiShampine5()
    AA, bb, cc = coefficients_as_in_tableau(A, b, c, tab)
    @test AA ≈ tab.A
    @test bb ≈ tab.α
    @test cc ≈ tab.c
  end

  let erk = Vern7()
    A, b, c = deduce_Butcher_tableau(erk)
    tab = constructVerner7()
    AA, bb, cc = coefficients_as_in_tableau(A, b, c, tab)
    @test AA ≈ tab.A
    @test bb ≈ tab.α
    @test cc ≈ tab.c
  end

  let erk = Vern8()
    A, b, c = deduce_Butcher_tableau(erk)
    tab = constructVerner8()
    AA, bb, cc = coefficients_as_in_tableau(A, b, c, tab)
    @test AA ≈ tab.A
    @test bb ≈ tab.α
    @test cc ≈ tab.c
  end
end

ode_tableau_tests(Float32)
ode_tableau_tests(Float64)


using JuMP
using MathOptInterface

const MOI = MathOptInterface

struct Solution
    nodes::Vector{Node}
    lsol::Array{Float64,2}
    πsol::Vector{Float64}
end

struct Result
    status::MOI.TerminationStatusCode
    sol::Solution
    value::Float64
end

include("nlp.jl")
include("soc.jl")


# Constants
const limit_kernel_width_default = 10.0
const speed_of_light = 299792458.0 # m/s

# spectrum.jl:  utils for computing simple spectrum models
include("spectrum.jl")
export AbstractSpectrum, SimpleSpectrum, SimulatedSpectrum, ConvolvedSpectrum
export doppler_shifted_spectrum
#export continuum_model, absorption_lines, gaussian_convolution_kernel
#export std_gaussian_line, absorption_line

# convolution_kernels.jl:  utils for convolution kernels
include("convolution_kernels.jl")
export AbstractConvolutionKernel, GaussianConvolutionKernel, GaussianMixtureConvolutionKernel


using IntervalConstraintProgramming
using ValidatedNumerics

using FactCheck

facts("Utilities") do
    @fact IntervalConstraintProgramming.unify_variables([:a, :c], [:c, :b]) -->
        ([:a,:b,:c], [1,3], [3,2], [1,0,2], [0,2,1])
end

facts("Separators") do
    II = -100..100
    X = IntervalBox(II, II)
    S = @constraint x^2 + y^2 <= 1

    inner, outer = S(X)
    @fact inner --> (-1..1, -1..1)
    @fact outer --> (II, II)

    X = IntervalBox(-∞..∞, -∞..∞)
    inner, outer = S(X)
    @fact inner --> (-1..1, -1..1)
    @fact outer --> (-∞..∞, -∞..∞)
end

facts("pave") do
    S1a = @constraint x > 0
    S1b = @constraint y > 0

    S1 = S1a ∩ S1b
    paving = pave(S1, IntervalBox(-3..3, -3..3), 0.1)
    @fact paving.inner --> [IntervalBox(1.5..3, 0..3), IntervalBox(0..1.5, 0..3)]
    @fact isempty(paving.boundary) --> true

    S2 = S1a ∪ S1b
    paving = pave(S2, IntervalBox(-3..3, -3..3), 0.1)
    @fact paving.inner --> [IntervalBox(-3..0, 0..3), IntervalBox(0..3, -3..3)]
    @fact isempty(paving.boundary) --> true

    S3 = @constraint x^2 + y^2 <= 1
    X = IntervalBox(-Inf..Inf, -Inf..Inf)
    paving = pave(S3, X, 1)

    @fact paving.inner --> [IntervalBox(Interval(0.0, 0.5), Interval(0.0, 0.8660254037844386)),
                    IntervalBox(Interval(0.0, 0.5), Interval(-0.8660254037844386, 0.0)),
                    IntervalBox(Interval(-0.5, 0.0), Interval(0.0, 0.8660254037844386)),
                    IntervalBox(Interval(-0.5, 0.0), Interval(-0.8660254037844386, 0.0))]

    @fact paving.boundary --> [ IntervalBox(Interval(0.5, 1.0), Interval(0.0, 0.8660254037844387)),
                        IntervalBox(Interval(0.0, 0.5), Interval(0.8660254037844386, 1.0)),
                        IntervalBox(Interval(0.5, 1.0), Interval(-0.8660254037844387, 0.0)),
                        IntervalBox(Interval(0.0, 0.5), Interval(-1.0, -0.8660254037844386)),
                        IntervalBox(Interval(-0.5, 0.0), Interval(0.8660254037844386, 1.0)),
                        IntervalBox(Interval(-1.0, -0.5), Interval(0.0, 0.8660254037844387)),
                        IntervalBox(Interval(-0.5, 0.0), Interval(-1.0, -0.8660254037844386)),
                        IntervalBox(Interval(-1.0, -0.5), Interval(-0.8660254037844387, 0.0))]
end

facts("Constants") do
    a = 3
    x = y = -∞..∞
    X = IntervalBox(x, y)
    S4 = @constraint x^2 + y^2 <= $a
    paving = pave(S4, X)

    @fact paving.ϵ --> 0.01
    @fact length(paving.inner) --> 1532
    @fact length(paving.boundary) --> 1536
end

facts("Volume") do
    x = 3..5
    @fact Vol(x).bounds --> 2

    V = Vol(IntervalBox(-1..1.5, 2..3.5))
    @fact typeof(V) --> IntervalConstraintProgramming.Vol{2, Float64}
    @fact V.bounds --> 3.75

end


struct BatchScalMvNormal{Tm<:AbstractMatrix,Tσ<:AbstractVector} <: ContinuousMatrixDistribution
    μ::Tm
    σ::Tσ
end

struct BatchDiagMvNormal{Tm<:AbstractMatrix,Tσ<:AbstractMatrix} <: ContinuousMatrixDistribution
    μ::Tm
    σ::Tσ
end

const BMN = Union{BatchScalMvNormal,BatchDiagMvNormal}

BatchMvNormal(μ::AbstractMatrix{T}, σ::AbstractVector{T}) where T<:Real = BatchScalMvNormal(μ,σ)
BatchMvNormal(μ::AbstractMatrix{T}, σ::AbstractMatrix{T}) where T<:Real = BatchDiagMvNormal(μ,σ)

#Base.length(d::BatchMvNormal) = size(d.μ,1)
Base.eltype(d::BMN) = eltype(d.μ)
Distributions.params(d::BMN) = (d.μ, d.σ)
Distributions.mean(d::BMN) = d.μ
Distributions.var(d::BatchDiagMvNormal) = d.σ .^2
Distributions.var(d::BatchScalMvNormal) = fillsimilar(d.σ,size(d.μ,1),1) .* reshape(d.σ .^2,1,:)

function Distributions.rand(d::BatchDiagMvNormal)
    μ, σ = d.μ, d.σ
    r = DistributionsAD.adapt_randn(Random.GLOBAL_RNG, μ, size(μ)...)
    μ .+ σ .* r
end

function Distributions.rand(d::BatchScalMvNormal)
    μ, σ = d.μ, reshape(d.σ, 1, :)
    r = DistributionsAD.adapt_randn(Random.GLOBAL_RNG, μ, size(μ)...)
    μ .+ σ .* r
end

function Distributions.logpdf(d::BMN, x::AbstractMatrix{T}) where T<:Real
    n = size(d.μ,1)
    μ = mean(d)
    σ2 = var(d)
    -(vec(sum(((x - μ).^2) ./ σ2 .+ log.(σ2), dims=1)) .+ n*log(T(2π))) ./ 2
end


# This is a Julia version of Solution of the optimal control problem
# based on code written by Andrea Walther. See:
# Walther, Andrea, and Narayanan, Sri Hari Krishna. Extending the Binomial Checkpointing
# Technique for Resilience. United States: N. p., 2016. https://www.osti.gov/biblio/1364654.

using Checkpointing
using ReverseDiff

include("optcontrolfunc.jl")

function header()
        println("**************************************************************************")
        println("*              Solution of the optimal control problem                   *")
        println("*                                                                        *")
        println("*                     J(y) = y_2(1) -> min                               *")
        println("*           s.t.   dy_1/dt = 0.5*y_1(t) + u(t),            y_1(0)=1      *")
        println("*                  dy_2/dt = y_1(t)^2 + 0.5*u(t)^2         y_2(0)=0      *")
        println("*                                                                        *")
        println("*                  the adjoints equations fulfill                        *")
        println("*                                                                        *")
        println("*         dl_1/dt = -0.5*l_1(t) - 2*y_1(t)*l_2(t)          l_1(1)=0      *")
        println("*         dl_2/dt = 0                                      l_2(1)=1      *")
        println("*                                                                        *")
        println("*   with Revolve for Online and (Multi-Stage) Offline Checkpointing      *")
        println("*                                                                        *")
        println("**************************************************************************")

        println("**************************************************************************")
        println("*        The solution of the optimal control problem above is            *")
        println("*                                                                        *")
        println("*        y_1*(t) = (2*e^(3t)+e^3)/(e^(3t/2)*(2+e^3))                     *")
        println("*        y_2*(t) = (2*e^(3t)-e^(6-3t)-2+e^6)/((2+e^3)^2)                 *")
        println("*          u*(t) = (2*e^(3t)-e^3)/(e^(3t/2)*(2+e^3))                     *")
        println("*        l_1*(t) = (2*e^(3-t)-2*e^(2t))/(e^(t/2)*(2+e^3))                *")
        println("*        l_2*(t) = 1                                                     *")
        println("*                                                                        *")
        println("**************************************************************************")

        return
end

function optcontrol(scheme, steps, adtool=ReverseDiffADTool())
    println( "\n STEPS    -> number of time steps to perform")
    println("SNAPS    -> number of checkpoints")
    println("INFO = 1 -> calculate only approximate solution")
    println("INFO = 2 -> calculate approximate solution + takeshots")
    println("INFO = 3 -> calculate approximate solution + all information ")
    println(" ENTER:   STEPS, SNAPS, INFO \n")

    h = 1.0/steps
    L = Array{Float64, 1}(undef, 2)
    L_H = Array{Float64, 1}(undef, 2)

    t = 0.0
    F = [1.0, 0.0]
    F_H = [0.0, 0.0]

    @checkpoint scheme adtool for i in 1:steps
        F_H = [F[1], F[2]]
        F = advance(F_H,t,h)
        t += h
    end

    F_opt = Array{Float64, 1}(undef, 2)
    L_opt = Array{Float64, 1}(undef, 2)
    opt_sol(F_opt,1.0)
    opt_lambda(L_opt,0.0)
    println("\n\n")
    println("y_1*(1)  = " , F_opt[1] , " y_2*(1)  = " , F_opt[2])
    println("y_1 (1)  = " , F[1] , "  y_2 (1)  = " , F[2] , " \n\n")
    println("l_1*(0)  = " , L_opt[1] , "  l_2*(0)  = " , L_opt[2])
    println("l_1 (0)  = " , L[1]     , "  sl_2 (0)  = " , L[2] , " ")
    return F_opt, F, L_opt, L
end

# main(10,3,3)


@printf("\nTesting linear operators calculus rules\n")

##########################
##### test Compose #######
##########################
m1, m2, m3 = 4, 7, 3
A1 = randn(m2, m1)
A2 = randn(m3, m2)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)

opC = Compose(opA2,opA1)
x = randn(m1)
y1 = test_op(opC, x, randn(m3), verb)
y2 = A2*A1*x
@test all(norm.(y1 .- y2) .<= 1e-12)

# test Compose longer
m1, m2, m3, m4 = 4, 7, 3, 2
A1 = randn(m2, m1)
A2 = randn(m3, m2)
A3 = randn(m4, m3)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opA3 = MatrixOp(A3)

opC1 = Compose(opA3,Compose(opA2,opA1))
opC2 = Compose(Compose(opA3,opA2),opA1)
x = randn(m1)
y1 = test_op(opC1, x, randn(m4), verb)
y2 = test_op(opC2, x, randn(m4), verb)
y3 = A3*A2*A1*x
@test all(norm.(y1 .- y2) .<= 1e-12)
@test all(norm.(y3 .- y2) .<= 1e-12)

#test Compose special cases
@test typeof(opA1*Eye(m1)) == typeof(opA1)
@test typeof(Eye(m2)*opA1) == typeof(opA1)
@test typeof(Eye(m2)*Eye(m2)) == typeof(Eye(m2))

opS1 = Compose(opA2,opA1)
opS1c = Scale(pi,opS1)
@test typeof(opS1c.A[end]) <: Scale

#properties
@test is_sliced(opC)            == false
@test is_linear(opC1)           == true
@test is_null(opC1)             == false
@test is_eye(opC1)              == false
@test is_diagonal(opC1)         == false
@test is_AcA_diagonal(opC1)     == false
@test is_AAc_diagonal(opC1)     == false
@test is_orthogonal(opC1)       == false
@test is_invertible(opC1)       == false
@test is_full_row_rank(opC1)    == false
@test is_full_column_rank(opC1) == false

# properties special case
opC = DCT((5,))*GetIndex((10,), 1:5)
@test is_sliced(opC)           == true
@test is_AAc_diagonal(opC)     == true
@test diag_AAc(opC)            == 1.0

d = randn(5)
opC = DiagOp(d)*GetIndex((10,), 1:5)
@test is_sliced(opC)           == true
@test is_diagonal(opC)         == true
@test diag(opC)                == d

# displacement test
m1, m2, m3, m4, m5 = 4, 7, 3, 2, 11
A1 = randn(m2, m1)
A2 = randn(m3, m2)
A3 = randn(m4, m3)
A4 = randn(m5, m4)
d1 = randn(m2)
d2 = pi
d3 = 0.0
d4 = randn(m5)
opA1 = AffineAdd(MatrixOp(A1),d1)
opA2 = AffineAdd(MatrixOp(A2),d2)
opA3 = MatrixOp(A3)
opA4 = AffineAdd(MatrixOp(A4),d4,false)

opC  = Compose(Compose(Compose(opA4,opA3),opA2),opA1)

x = randn(m1)

@test norm( opC*x - (A4*(A3*( A2*( A1*x+d1 ) .+ d2 ) .+ d3)-d4) )  < 1e-9
@test norm( displacement(opC) - ( A4*(A3*(A2*d1 .+d2 ) .+ d3)-d4) ) < 1e-9

opA4 = MatrixOp(A4)
opC  = AffineAdd(Compose(Compose(Compose(opA4,opA3),opA2),opA1),d4)
@test norm( opC*x - (A4*(A3*( A2*( A1*x+d1 ) .+ d2 ) .+ d3)+d4) )  < 1e-9
@test norm( displacement(opC) - ( A4*(A3*(A2*d1 .+ d2) .+ d3)+d4) ) < 1e-9

@test norm( remove_displacement(opC)*x - (A4*(A3*( A2*( A1*x ) ) ) ) )  < 1e-9

########################
### test DCAT    #######
########################

m1, n1, m2, n2 = 4, 7, 5, 2
A1 = randn(m1, n1)
A2 = randn(m2, n2)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opD = DCAT(opA1, opA2)
x1 = randn(n1)
x2 = randn(n2)
y1 = test_op(opD, ArrayPartition(x1, x2), ArrayPartition(randn(m1),randn(m2)), verb)
y2 = ArrayPartition(A1*x1, A2*x2)
@test norm(y1 .- y2) .<= 1e-12

# test DCAT longer

m1, n1, m2, n2, m3, n3 = 4, 7, 5, 2, 5, 5
A1 = randn(m1, n1)
A2 = randn(m2, n2)
A3 = randn(m3, n3)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opA3 = MatrixOp(A3)
opD = DCAT(opA1, opA2, opA3)
x1 = randn(n1)
x2 = randn(n2)
x3 = randn(n3)
y1 = test_op(opD, ArrayPartition(x1, x2, x3), ArrayPartition(randn(m1),randn(m2),randn(m3)), verb)
y2 = ArrayPartition(A1*x1, A2*x2, A3*x3)
@test norm(y1 .- y2) .<= 1e-12

#properties
@test is_linear(opD)           == true
@test is_null(opD)             == false
@test is_eye(opD)              == false
@test is_diagonal(opD)         == false
@test is_AcA_diagonal(opD)     == false
@test is_AAc_diagonal(opD)     == false
@test is_orthogonal(opD)       == false
@test is_invertible(opD)       == false
@test is_full_row_rank(opD)    == false
@test is_full_column_rank(opD) == false

# DCAT of Eye

n1, n2 = 4, 7
x1 = randn(n1)
x2 = randn(n2)

opD = Eye(ArrayPartition(x1,x2))
y1 = test_op(opD, ArrayPartition(x1, x2), ArrayPartition(randn(n1),randn(n2)), verb)

#properties
@test is_linear(opD)           == true
@test is_null(opD)             == false
@test is_eye(opD)              == true
@test is_diagonal(opD)         == true
@test is_AcA_diagonal(opD)     == true
@test is_AAc_diagonal(opD)     == true
@test is_orthogonal(opD)       == true
@test is_invertible(opD)       == true
@test is_full_row_rank(opD)    == true
@test is_full_column_rank(opD) == true

@test diag(opD) == 1
@test diag_AcA(opD) == 1
@test diag_AAc(opD) == 1

# displacement DCAT

m1, n1, m2, n2, m3, n3 = 4, 7, 5, 2, 5, 5
A1 = randn(m1, n1)
A2 = randn(m2, n2)
A3 = randn(m3, n3)
d1 = randn(m1)
d2 = randn(m2)
d3 = randn(m3)
opA1 = AffineAdd(MatrixOp(A1), d1)
opA2 = AffineAdd(MatrixOp(A2), d2)
opA3 = AffineAdd(MatrixOp(A3), d3)
opD = DCAT(opA1, opA2, opA3)
x1 = randn(n1)
x2 = randn(n2)
x3 = randn(n3)
y1 = opD*ArrayPartition(x1,x2,x3)
y2 = ArrayPartition(A1*x1+d1, A2*x2+d2, A3*x3+d3)
@test norm(y1 .- y2) .<= 1e-12
@test norm(displacement(opD) .- ArrayPartition(d1,d2,d3)) .<= 1e-12
y1 = remove_displacement(opD)*ArrayPartition(x1,x2,x3)
y2 = ArrayPartition(A1*x1, A2*x2, A3*x3)
@test norm(y1 .- y2) .<= 1e-12

#######################
## test HCAT    #######
#######################

m, n1, n2 = 4, 7, 5
A1 = randn(m, n1)
A2 = randn(m, n2)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opH = HCAT(opA1, opA2)
x1 = randn(n1)
x2 = randn(n2)
y1 = test_op(opH, ArrayPartition(x1, x2), randn(m), verb)
y2 = A1*x1 + A2*x2
@test norm(y1-y2) <= 1e-12

#permuatation
p = [2;1]
opHp = opH[p]
y1 = test_op(opHp, ArrayPartition(x2, x1), randn(m), verb)
@test norm(y1-y2) <= 1e-12

# test HCAT longer

m, n1, n2, n3 = 4, 7, 5, 6
A1 = randn(m, n1)
A2 = randn(m, n2)
A3 = randn(m, n3)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opA3 = MatrixOp(A3)
opH = HCAT(opA1, opA2, opA3)
x1 = randn(n1)
x2 = randn(n2)
x3 = randn(n3)
y1 = test_op(opH, ArrayPartition(x1, x2, x3), randn(m), verb)
y2 = A1*x1 + A2*x2 + A3*x3
@test norm(y1-y2) <= 1e-12

# test HCAT of HCAT
opHH = HCAT(opH, opA2, opA3)
y1 = test_op(opHH, ArrayPartition(x1, x2, x3, x2, x3), randn(m), verb)
y2 = A1*x1 + A2*x2 + A3*x3 + A2*x2 + A3*x3
@test norm(y1-y2) <= 1e-12

opHH = HCAT(opH, opH, opA3)
x = ArrayPartition(x1, x2, x3, x1, x2, x3, x3)
y1 = test_op(opHH, x, randn(m), verb)
y2 = A1*x1 + A2*x2 + A3*x3 + A1*x1 + A2*x2 + A3*x3 + A3*x3
@test norm(y1-y2) <= 1e-12

opA3 = MatrixOp(randn(n1,n1))
@test_throws Exception HCAT(opA1,opA2,opA3)
opF = AbstractOperators.DFT(Complex{Float64},(m,))
@test_throws Exception HCAT(opA1,opF,opA2)

# test utilities

# permutation
p = randperm(ndoms(opHH,2))
opHP = AbstractOperators.permute(opHH,p)

xp = ArrayPartition(x.x[p]...)

y1 = test_op(opHP, xp, randn(m), verb)

pp = randperm(ndoms(opHH,2))
opHPP = AbstractOperators.permute(opHH,pp)
xpp = ArrayPartition(x.x[pp]...)
y1 = test_op(opHPP, xpp, randn(m), verb)

#properties
m, n1, n2, n3 = 4, 7, 5, 6
A1 = randn(m, n1)
A2 = randn(m, n2)
A3 = randn(m, n3)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opA3 = MatrixOp(A3)
op = HCAT(opA1, opA2, opA3)
@test is_linear(op)           == true
@test is_null(op)             == false
@test is_eye(op)              == false
@test is_diagonal(op)         == false
@test is_AcA_diagonal(op)     == false
@test is_AAc_diagonal(op)     == false
@test is_orthogonal(op)       == false
@test is_invertible(op)       == false
@test is_full_row_rank(op)    == true
@test is_full_column_rank(op) == false

d = randn(n1).+im.*randn(n1)
op = HCAT(DiagOp(d), AbstractOperators.DFT(Complex{Float64},n1))
@test is_null(op)             == false
@test is_eye(op)              == false
@test is_diagonal(op)         == false
@test is_AcA_diagonal(op)     == false
@test is_AAc_diagonal(op)     == true
@test is_orthogonal(op)       == false
@test is_invertible(op)       == false
@test is_full_row_rank(op)    == true
@test is_full_column_rank(op) == false

@test diag_AAc(op) == d .* conj(d) .+ n1

y1 = randn(n1) .+ im .* randn(n1)
@test norm(op*(op'*y1).-diag_AAc(op).*y1) <1e-12

#test displacement

m, n1, n2 = 4, 7, 5
A1 = randn(m, n1)
A2 = randn(m, n2)
d1 = randn(m)
d2 = randn(m)
opA1 = AffineAdd(MatrixOp(A1), d1)
opA2 = AffineAdd(MatrixOp(A2), d2)
opH = HCAT(opA1, opA2)
x1 = randn(n1)
x2 = randn(n2)
y1 = opH*ArrayPartition(x1,x2)
y2 = A1*x1+d1 + A2*x2+d2
@test norm(y1-y2) <= 1e-12
y1 = remove_displacement(opH)*ArrayPartition(x1,x2)
y2 = A1*x1 + A2*x2
@test norm(y1-y2) <= 1e-12

########################
### test Reshape #######
########################

m, n = 8, 4
dim_out = (2, 2, 2)
A1 = randn(m, n)
opA1 = MatrixOp(A1)
opR = Reshape(opA1, dim_out)
opR = Reshape(opA1, dim_out...)
x1 = randn(n)
y1 = test_op(opR, x1, randn(dim_out), verb)
y2 = reshape(A1*x1, dim_out)
@test norm(y1-y2) <= 1e-12

@test_throws Exception Reshape(opA1,(2,2,1))

@test is_null(opR)             == is_null(opA1)
@test is_eye(opR)              == is_eye(opA1)
@test is_diagonal(opR)         == is_diagonal(opA1)
@test is_AcA_diagonal(opR)     == is_AcA_diagonal(opA1)
@test is_AAc_diagonal(opR)     == is_AAc_diagonal(opA1)
@test is_orthogonal(opR)       == is_orthogonal(opA1)
@test is_invertible(opR)       == is_invertible(opA1)
@test is_full_row_rank(opR)    == is_full_row_rank(opA1)
@test is_full_column_rank(opR) == is_full_column_rank(opA1)

# testing displacement
m, n = 8, 4
dim_out = (2, 2, 2)
A1 = randn(m, n)
d1 = randn(m)
opA1 = AffineAdd(MatrixOp(A1),d1)
opR = Reshape(opA1, dim_out)
x1 = randn(n)
y1 = opR*x1
y2 = reshape(A1*x1+d1, dim_out)
@test norm(y1-y2) <= 1e-12
y1 = remove_displacement(opR)*x1
y2 = reshape(A1*x1, dim_out)
@test norm(y1-y2) <= 1e-12

#######################
## test Scale   #######
#######################

m, n = 8, 4
coeff = pi
A1 = randn(m, n)
opA1 = MatrixOp(A1)
opS = Scale(coeff, opA1)
x1 = randn(n)
y1 = test_op(opS, x1, randn(m), verb)
y2 = coeff*A1*x1
@test norm(y1-y2) <= 1e-12

coeff2 = 3
opS2 = Scale(coeff2, opS)
y1 = test_op(opS2, x1, randn(m), verb)
y2 = coeff2*coeff*A1*x1
@test norm(y1-y2) <= 1e-12

opF = AbstractOperators.DFT(m,n)
opS = Scale(coeff, opF)
x1 = randn(m,n)
y1 = test_op(opS, x1, fft(randn(m,n)), verb)
y2 = coeff*(fft(x1))
@test norm(y1-y2) <= 1e-12

opS = Scale(coeff, opA1)
@test is_null(opS)             == is_null(opA1)
@test is_eye(opS)              == is_eye(opA1)
@test is_diagonal(opS)         == is_diagonal(opA1)
@test is_AcA_diagonal(opS)     == is_AcA_diagonal(opA1)
@test is_AAc_diagonal(opS)     == is_AAc_diagonal(opA1)
@test is_orthogonal(opS)       == is_orthogonal(opA1)
@test is_invertible(opS)       == is_invertible(opA1)
@test is_full_row_rank(opS)    == is_full_row_rank(opA1)
@test is_full_column_rank(opS) == is_full_column_rank(opA1)

op = Scale(-4.0,AbstractOperators.DFT(10))
@test is_AAc_diagonal(op)     == true
@test diag_AAc(op) == 16*10

op = Scale(-4.0,ZeroPad((10,), 20))
@test is_AcA_diagonal(op)     == true
@test diag_AcA(op) == 16

# special case, Scale of DiagOp gets a DiagOp
d = randn(10)
op = Scale(3,DiagOp(d))
@test typeof(op) <: DiagOp
@test norm(diag(op) - 3 .*d) < 1e-12

# Scale with imaginary coeff gives error
m, n = 8, 4
coeff = im
A1 = randn(m, n)
opA1 = MatrixOp(A1)
@test_throws ErrorException Scale(coeff, opA1)

## testing displacement
m, n = 8, 4
coeff = pi
A1 = randn(m, n)
d1 = randn(m)
opA1 = AffineAdd(MatrixOp(A1),d1)
opS = Scale(coeff, opA1)
x1 = randn(n)
y1 = opS*x1
y2 = coeff*(A1*x1+d1)
@test norm(y1-y2) <= 1e-12
y1 = remove_displacement(opS)*x1
y2 = coeff*(A1*x1)
@test norm(y1-y2) <= 1e-12

########################
### test Sum     #######
########################

m,n = 5,7
A1 = randn(m,n)
A2 = randn(m,n)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opS = Sum(opA1,opA2)
x1 = randn(n)
y1 = test_op(opS, x1, randn(m), verb)
y2 = A1*x1+A2*x1
@test norm(y1-y2) <= 1e-12

#test Sum longer
m,n = 5,7
A1 = randn(m,n)
A2 = randn(m,n)
A3 = randn(m,n)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opA3 = MatrixOp(A3)
opS = Sum(opA1,opA2,opA3)
x1 = randn(n)
y1 = test_op(opS, x1, randn(m), verb)
y2 = A1*x1+A2*x1+A3*x1
@test norm(y1-y2) <= 1e-12

opA3 = MatrixOp(randn(m,m))
@test_throws Exception Sum(opA1,opA3)
opF = AbstractOperators.DFT(Float64,(m,))
@test_throws Exception Sum(opF,opA3)

@test is_null(opS)             == false
@test is_eye(opS)              == false
@test is_diagonal(opS)         == false
@test is_AcA_diagonal(opS)     == false
@test is_AAc_diagonal(opS)     == false
@test is_orthogonal(opS)       == false
@test is_invertible(opS)       == false
@test is_full_row_rank(opS)    == true
@test is_full_column_rank(opS) == false

d = randn(10)
op = Sum(Scale(-3.1,Eye(10)),DiagOp(d))
@test is_diagonal(op)         == true
@test norm(   diag(op) - (d .-3.1)  )<1e-12

#test displacement of sum
m,n = 5,7
A1 = randn(m,n)
A2 = randn(m,n)
A3 = randn(m,n)
d1 = randn(m)
d2 = pi
d3 = randn(m)
opA1 = AffineAdd(MatrixOp(A1), d1)
opA2 = AffineAdd(MatrixOp(A2), d2)
opA3 = AffineAdd(MatrixOp(A3), d3)
opS = Sum(opA1,opA2,opA3)
x1 = randn(n)
y2 = A1*x1+A2*x1+A3*x1+d1.+d2+d3
@test norm(opS*x1-y2) <= 1e-12
@test norm(displacement(opS) - (d1.+d2+d3)) <= 1e-12
y2 = A1*x1+A2*x1+A3*x1
@test norm(remove_displacement(opS)*x1-y2) <= 1e-12

###################################
###### test AdjointOperator #######
###################################

m,n = 5,7
A1 = randn(m,n)
opA1 = MatrixOp(A1)
opA1t = MatrixOp(A1')
opT = AdjointOperator(opA1)
x1 = randn(m)
y1 = test_op(opT, x1, randn(n), verb)
y2 = A1'*x1
@test norm(y1-y2) <= 1e-12

@test is_null(opT)             == is_null(opA1t)
@test is_eye(opT)              == is_eye(opA1t)
@test is_diagonal(opT)         == is_diagonal(opA1t)
@test is_AcA_diagonal(opT)     == is_AcA_diagonal(opA1t)
@test is_AAc_diagonal(opT)     == is_AAc_diagonal(opA1t)
@test is_orthogonal(opT)       == is_orthogonal(opA1t)
@test is_invertible(opT)       == is_invertible(opA1t)
@test is_full_row_rank(opT)    == is_full_row_rank(opA1t)
@test is_full_column_rank(opT) == is_full_column_rank(opA1t)

d = randn(3)
op = AdjointOperator(DiagOp(d))
@test is_diagonal(op) == true
@test diag(op) == d

op = AdjointOperator(ZeroPad((10,),5))
@test is_AcA_diagonal(op) == false
@test is_AAc_diagonal(op) == true
@test diag_AAc(op) == 1

#############################
######## test VCAT    #######
#############################

m1, m2, n = 4, 7, 5
A1 = randn(m1, n)
A2 = randn(m2, n)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opV = VCAT(opA1, opA2)
x1 = randn(n)
y1 = test_op(opV, x1, ArrayPartition(randn(m1), randn(m2)), verb)
y2 = ArrayPartition(A1*x1, A2*x1)
@test norm(y1 - y2) .<= 1e-12

m1, n = 4, 5
A1 = randn(m1, n)+im*randn(m1, n)
opA1 = MatrixOp(A1)
opA2 = AbstractOperators.DFT(n)'
opV = VCAT(opA1, opA2)
x1 = randn(n)+im*randn(n)
y1 = test_op(opV, x1, ArrayPartition(randn(m1)+im*randn(m1), randn(n)), verb)
y2 = ArrayPartition(A1*x1, opA2*x1)
@test norm(y1 - y2) .<= 1e-12

#test VCAT longer
m1, m2, m3, n = 4, 7, 3, 5
A1 = randn(m1, n)
A2 = randn(m2, n)
A3 = randn(m3, n)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opA3 = MatrixOp(A3)
opV = VCAT(opA1, opA2, opA3)
x1 = randn(n)
y1 = test_op(opV, x1, ArrayPartition(randn(m1), randn(m2), randn(m3)), verb)
y2 = ArrayPartition(A1*x1, A2*x1, A3*x1)
@test norm(y1 - y2) .<= 1e-12

#test VCAT of VCAT
opVV = VCAT(opV,opA3)
y1 = test_op(opVV, x1, ArrayPartition(randn(m1), randn(m2), randn(m3), randn(m3)), verb)
y2 = ArrayPartition(A1*x1, A2*x1, A3*x1, A3*x1)
@test norm(y1 .- y2) <= 1e-12

opVV = VCAT(opA1,opV,opA3)
y1 = test_op(opVV, x1, ArrayPartition(randn(m1), randn(m1), randn(m2), randn(m3), randn(m3)), verb)
y2 = ArrayPartition(A1*x1, A1*x1, A2*x1, A3*x1, A3*x1)
@test norm(y1 - y2) <= 1e-12

opA3 = MatrixOp(randn(m1,m1))
@test_throws Exception VCAT(opA1,opA2,opA3)
opF = AbstractOperators.DFT(Complex{Float64},(n,))
@test_throws Exception VCAT(opA1,opF,opA2)

###properties
m1, m2, m3, n = 4, 7, 3, 5
A1 = randn(m1, n)
A2 = randn(m2, n)
A3 = randn(m3, n)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opA3 = MatrixOp(A3)
op = VCAT(opA1, opA2, opA3)
@test is_linear(op)           == true
@test is_null(op)             == false
@test is_eye(op)              == false
@test is_diagonal(op)         == false
@test is_AcA_diagonal(op)     == false
@test is_AAc_diagonal(op)     == false
@test is_orthogonal(op)       == false
@test is_invertible(op)       == false
@test is_full_row_rank(op)    == false
@test is_full_column_rank(op) == true

op = VCAT(AbstractOperators.DFT(Complex{Float64},10), Eye(Complex{Float64},10) )
@test is_AcA_diagonal(op)     == true
@test diag_AcA(op) == 11

##test displacement
m1, m2, n = 4, 7, 5
A1 = randn(m1, n)
A2 = randn(m2, n)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
d1 = randn(m1)
d2 = randn(m2)
opV = VCAT(AffineAdd(opA1,d1), AffineAdd(opA2,d2))
x1 = randn(n)
y1 = opV*x1
y2 = ArrayPartition(A1*x1+d1, A2*x1+d2)
@test norm(y1 - y2) <= 1e-12
y1 = remove_displacement(opV)*x1
y2 = ArrayPartition(A1*x1, A2*x1)
@test norm(y1 - y2) <= 1e-12

#########################
#### test BroadCast #####
#########################

m, n = 8, 4
dim_out = (m, 10)
A1 = randn(m, n)
opA1 = MatrixOp(A1)
opR = BroadCast(opA1, dim_out)
x1 = randn(n)
y1 = test_op(opR, x1, randn(dim_out), verb)
y2 = zeros(dim_out)
y2 .= A1*x1
@test norm(y1-y2) <= 1e-12

m, n, l, k = 8, 4, 5, 7
dim_out = (m, n, l, k)
opA1 = Eye(m,n)
opR = BroadCast(opA1, dim_out)
x1 = randn(m,n)
y1 = test_op(opR, x1, randn(dim_out), verb)
y2 = zeros(dim_out)
y2 .= x1
@test norm(y1-y2) <= 1e-12

@test_throws Exception BroadCast(opA1,(m,m))

m, l = 1, 5
dim_out = (m, l)
opA1 = Eye(m)
opR = BroadCast(opA1, dim_out)
x1 = randn(m)
y1 = test_op(opR, x1, randn(dim_out), verb)
y2 = zeros(dim_out)
y2 .= x1
@test norm(y1-y2) <= 1e-12

#colum in - matrix out
m, l = 4, 5
dim_out = (m, l)
opA1 = Eye(1,l)
opR = BroadCast(opA1, dim_out)
x1 = randn(1,l)
y1 = test_op(opR, x1, randn(dim_out), verb)
y2 = zeros(dim_out)
y2 .= x1
@test norm(y1-y2) <= 1e-12

op = HCAT(Eye(m,l),opR)
x1 = ArrayPartition(randn(m,l),randn(1,l))
y1 = test_op(op, x1, randn(dim_out), verb)
y2 = x1.x[1].+x1.x[2]
@test norm(y1-y2) <= 1e-12

m, n, l  = 2, 5, 8
dim_out = (m, n, l)
opA1 = Eye(m)
opR = BroadCast(opA1, dim_out)
x1 = randn(m)
y1 = test_op(opR, x1, randn(dim_out), verb)
y2 = zeros(dim_out)
y2 .= x1
@test norm(y1-y2) <= 1e-12

m, n, l  = 1, 5, 8
dim_out = (m, n, l)
opA1 = Eye(m)
opR = BroadCast(opA1, dim_out)
x1 = randn(m)
y1 = test_op(opR, x1, randn(dim_out), verb)
y2 = zeros(dim_out)
y2 .= x1
@test norm(y1-y2) <= 1e-12

m, n, l  = 1, 5, 8
dim_out = (m, n, l)
opA1 = Scale(2.4,Eye(m))
opR = BroadCast(opA1, dim_out)
x1 = randn(m)
y1 = test_op(opR, x1, randn(dim_out), verb)
y2 = zeros(dim_out)
y2 .= 2.4*x1
@test norm(y1-y2) <= 1e-12

@test is_null(opR)             == is_null(opA1)
@test is_eye(opR)              == false
@test is_diagonal(opR)         == false
@test is_AcA_diagonal(opR)     == false
@test is_AAc_diagonal(opR)     == false
@test is_orthogonal(opR)       == false
@test is_invertible(opR)       == false
@test is_full_row_rank(opR)    == false
@test is_full_column_rank(opR) == false

# test displacement

m, n = 8, 4
dim_out = (m, 10)
A1 = randn(m, n)
d1 = randn(m)
opA1 = AffineAdd(MatrixOp(A1), d1)
opR = BroadCast(opA1, dim_out)
x1 = randn(n)
y1 = opR*x1
y2 = zeros(dim_out)
y2 .= A1*x1+d1
@test norm(y1-y2) <= 1e-12
x1 = randn(n)
y1 = remove_displacement(opR)*x1
y2 = zeros(dim_out)
y2 .= A1*x1
@test norm(y1-y2) <= 1e-12

###########################
##### test AffineAdd  #####
###########################

n,m = 5,6
A = randn(n,m)
opA = MatrixOp(A)
d = randn(n)
T = AffineAdd(opA,d)

println(T)
x1 = randn(m)
y1 = T*x1
@test norm(y1-(A*x1+d)) <1e-9
r = randn(n)
@test norm(T'*r-(A'*r)) <1e-9
@test displacement(T) == d
@test norm(remove_displacement(T)*x1-A*x1) <1e-9

# with sign
T = AffineAdd(opA,d,false)
@test sign(T) == -1

println(T)
x1 = randn(m)
y1 = T*x1
@test norm(y1-(A*x1-d)) <1e-9
r = randn(n)
@test norm(T'*r-(A'*r)) <1e-9
@test displacement(T) == -d
@test norm(remove_displacement(T)*x1-A*x1) <1e-9

# with scalar
T = AffineAdd(opA,pi)
@test sign(T) == 1

println(T)
x1 = randn(m)
y1 = T*x1
@test norm(y1-(A*x1.+pi)) <1e-9
r = randn(n)
@test norm(T'*r-(A'*r)) < 1e-9
@test displacement(T) .- pi < 1e-9
@test norm(remove_displacement(T)*x1-A*x1) <1e-9

@test_throws DimensionMismatch AffineAdd(MatrixOp(randn(2,5)),randn(5))
@test_throws ErrorException AffineAdd(AbstractOperators.DFT(4),randn(4))
AffineAdd(AbstractOperators.DFT(4),pi)
@test_throws ErrorException AffineAdd(Eye(4),im*pi)

# with scalar and vector
d = randn(n)
T = AffineAdd(AffineAdd(opA,pi),d,false)

println(T)
x1 = randn(m)
y1 = T*x1
@test norm(y1-(A*x1 .+ pi .- d)) <1e-9
r = randn(n)
@test norm(T'*r-(A'*r)) < 1e-9
@test norm(displacement(T) .- (pi .-d )) < 1e-9

T2 = remove_displacement(T)
@test norm(T2*x1-(A*x1)) <1e-9

# permute AddAffine
n,m = 5,6
A = randn(n,m)
d = randn(n)
opH = HCAT(Eye(n),MatrixOp(A))
x = ArrayPartition(randn(n),randn(m))
opHT = AffineAdd(opH,d)

@test norm(opHT*x-(x.x[1]+A*x.x[2].+d)) < 1e-12
p = [2;1]
@test norm(AbstractOperators.permute(opHT,p)*ArrayPartition(x.x[p]...)-(x.x[1]+A*x.x[2].+d)) < 1e-12

############################
####### test combin. #######
############################

## test Compose of HCAT
m1, m2, m3, m4 = 4, 7, 3, 2
A1 = randn(m3, m1)
A2 = randn(m3, m2)
A3 = randn(m4, m3)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opA3 = MatrixOp(A3)
opH = HCAT(opA1,opA2)
opC = Compose(opA3,opH)
x1, x2 = randn(m1), randn(m2)
y1 = test_op(opC, ArrayPartition(x1,x2), randn(m4), verb)

y2 = A3*(A1*x1+A2*x2)

@test norm(y1-y2) < 1e-9

opCp = AbstractOperators.permute(opC,[2,1])
y1 = test_op(opCp, ArrayPartition(x2,x1), randn(m4), verb)

@test norm(y1-y2) < 1e-9

## test HCAT of Compose of HCAT
m5 = 10
A4 = randn(m4,m5)
x3 = randn(m5)
opHC = HCAT(opC,MatrixOp(A4))
x = ArrayPartition(x1,x2,x3)
y1 = test_op(opHC, x, randn(m4), verb)

@test norm(y1-(y2+A4*x3)) < 1e-9

p = randperm(ndoms(opHC,2))
opHP = AbstractOperators.permute(opHC,p)

xp = ArrayPartition(x.x[p]...)

y1 = test_op(opHP, xp, randn(m4), verb)

pp = randperm(ndoms(opHC,2))
opHPP = AbstractOperators.permute(opHC,pp)
xpp = ArrayPartition(x.x[pp]...)
y1 = test_op(opHPP, xpp, randn(m4), verb)

# test VCAT of HCAT's
m1, m2, n1 = 4, 7, 3
A1 = randn(n1, m1)
A2 = randn(n1, m2)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opH1 = HCAT(opA1,opA2)

m1, m2, n2 = 4, 7, 5
A3 = randn(n2, m1)
A4 = randn(n2, m2)
opA3 = MatrixOp(A3)
opA4 = MatrixOp(A4)
opH2 = HCAT(opA3,opA4)

opV = VCAT(opH1,opH2)
x1, x2 = randn(m1), randn(m2)
y1 = test_op(opV, ArrayPartition(x1,x2), ArrayPartition(randn(n1),randn(n2)), verb)
y2 = ArrayPartition(A1*x1+A2*x2,A3*x1+A4*x2)
@test norm(y1 - y2) <= 1e-12

# test VCAT of HCAT's with complex num
m1, m2, n1 = 4, 7, 5
A1 = randn(n1, m1)+im*randn(n1, m1)
opA1 = MatrixOp(A1)
opA2 = AbstractOperators.DFT(n1)
opH1 = HCAT(opA1,opA2)

m1, m2, n2 = 4, 7, 5
A3 = randn(n2, m1)+im*randn(n2,m1)
opA3 = MatrixOp(A3)
opA4 = AbstractOperators.DFT(n2)
opH2 = HCAT(opA3,opA4)

opV = VCAT(opH1,opH2)
x1, x2 = randn(m1)+im*randn(m1), randn(n2)
y1 = test_op(opV, ArrayPartition(x1,x2), ArrayPartition(randn(n1)+im*randn(n1),randn(n2)+im*randn(n2)), verb)
y2 = ArrayPartition(A1*x1+fft(x2),A3*x1+fft(x2))
@test norm(y1 - y2) <= 1e-12

# test HCAT of VCAT's

n1, n2, m1, m2 = 3, 5, 4, 7
A = randn(m1, n1); opA = MatrixOp(A)
B = randn(m1, n2); opB = MatrixOp(B)
C = randn(m2, n1); opC = MatrixOp(C)
D = randn(m2, n2); opD = MatrixOp(D)
opV = HCAT(VCAT(opA, opC), VCAT(opB, opD))
x1 = randn(n1)
x2 = randn(n2)
y1 = test_op(opV, ArrayPartition(x1, x2), ArrayPartition(randn(m1), randn(m2)), verb)
y2 = ArrayPartition(A*x1 + B*x2, C*x1 + D*x2)

@test norm(y1 - y2) <= 1e-12

# test Sum of HCAT's

m, n1, n2, n3 = 4, 7, 5, 3
A1 = randn(m, n1)
A2 = randn(m, n2)
A3 = randn(m, n3)
B1 = randn(m, n1)
B2 = randn(m, n2)
B3 = randn(m, n3)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opA3 = MatrixOp(A3)
opB1 = MatrixOp(B1)
opB2 = MatrixOp(B2)
opB3 = MatrixOp(B3)
opHA = HCAT(opA1, opA2, opA3)
opHB = HCAT(opB1, opB2, opB3)
opS = Sum(opHA, opHB)
x1 = randn(n1)
x2 = randn(n2)
x3 = randn(n3)
y1 = test_op(opS, ArrayPartition(x1, x2, x3), randn(m), verb)
y2 = A1*x1 + B1*x1 + A2*x2 + B2*x2 + A3*x3 + B3*x3

@test norm(y1-y2) <= 1e-12

p = [3;2;1]
opSp = AbstractOperators.permute(opS,p)
y1 = test_op(opSp, ArrayPartition(((x1, x2, x3)[p])...), randn(m), verb)

# test Sum of VCAT's

m1, m2, n = 4, 7, 5
A1 = randn(m1, n)
A2 = randn(m2, n)
B1 = randn(m1, n)
B2 = randn(m2, n)
C1 = randn(m1, n)
C2 = randn(m2, n)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opB1 = MatrixOp(B1)
opB2 = MatrixOp(B2)
opC1 = MatrixOp(C1)
opC2 = MatrixOp(C2)
opVA = VCAT(opA1, opA2)
opVB = VCAT(opB1, opB2)
opVC = VCAT(opC1, opC2)
opS = Sum(opVA, opVB, opVC)
x = randn(n)
y1 = test_op(opS, x, ArrayPartition(randn(m1), randn(m2)), verb)
y2 = ArrayPartition(A1*x + B1*x +C1*x, A2*x + B2*x + C2*x)

@test norm(y1 - y2) .<= 1e-12

# test Scale of DCAT

m1, n1 = 4, 7
m2, n2 = 3, 5
A1 = randn(m1, n1)
A2 = randn(m2, n2)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opD = DCAT(opA1, opA2)
coeff = randn()
opS = Scale(coeff, opD)
x1 = randn(n1)
x2 = randn(n2)
y = test_op(opS, ArrayPartition(x1, x2), ArrayPartition(randn(m1), randn(m2)), verb)
z = ArrayPartition(coeff*A1*x1, coeff*A2*x2)

@test norm(y - z) <= 1e-12

# test Scale of VCAT

m1, m2, n = 4, 3, 7
A1 = randn(m1, n)
A2 = randn(m2, n)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opV = VCAT(opA1, opA2)
coeff = randn()
opS = Scale(coeff, opV)
x = randn(n)
y = test_op(opS, x, ArrayPartition(randn(m1), randn(m2)), verb)
z = ArrayPartition(coeff*A1*x, coeff*A2*x)

@test norm(y - z) <= 1e-12

# test Scale of HCAT

m, n1, n2 = 4, 3, 7
A1 = randn(m, n1)
A2 = randn(m, n2)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opH = HCAT(opA1, opA2)
coeff = randn()
opS = Scale(coeff, opH)
x1 = randn(n1)
x2 = randn(n2)
y = test_op(opS, ArrayPartition(x1, x2), randn(m), verb)
z = coeff*(A1*x1 + A2*x2)

@test norm(y - z) <= 1e-12

# test DCAT of HCATs

m1, m2, n1, n2, l1, l2, l3 = 2,3,4,5,6,7,8
A1 = randn(m1, n1)
A2 = randn(m1, n2)
B1 = randn(m2, n1)
B2 = randn(m2, n2)
B3 = randn(m2, n2)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opH1 = HCAT(opA1, opA2)
opB1 = MatrixOp(B1)
opB2 = MatrixOp(B2)
opB3 = MatrixOp(B3)
opH2 = HCAT(opB1, opB2, opB3)

op = DCAT(opA1, opH2)
x =  ArrayPartition(randn.(size(op,2))...)
y0 = ArrayPartition(randn.(size(op,1))...)
y = test_op(op, x, y0, verb)

op = DCAT(opH1, opH2)
x =  ArrayPartition(randn.(size(op,2))...)
y0 = ArrayPartition(randn.(size(op,1))...)
y = test_op(op, x, y0, verb)

p = randperm(ndoms(op,2))
y2 = op[p]*ArrayPartition(x.x[p]...)

@test norm(y - y2) <= 1e-8

# test Scale of Sum

m,n = 5,7
A1 = randn(m,n)
A2 = randn(m,n)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
opS = Sum(opA1,opA2)
coeff = pi
opSS = Scale(coeff,opS)
x1 = randn(n)
y1 = test_op(opSS, x1, randn(m), verb)
y2 = coeff*(A1*x1+A2*x1)
@test norm(y1-y2) <= 1e-12

# test Scale of Compose

m1, m2, m3 = 4, 7, 3
A1 = randn(m2, m1)
A2 = randn(m3, m2)
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)

coeff = pi
opC = Compose(opA2,opA1)
opS = Scale(coeff,opC)
x = randn(m1)
y1 = test_op(opS, x, randn(m3), verb)
y2 = coeff*(A2*A1*x)
@test all(norm.(y1 .- y2) .<= 1e-12)


#
# Level Sets
#
#    Functions for computing level curves and surfaces in simplicial complexes
#
#    Author: Carlos Valero Valdés
#    e-mail: valerocar@gmail.com
#
#

"""
    levelcurve(ps::Array{<:Real,2}, fs::Array{<:Int,2}, vals::::Array{<:Real,1}, level::Real)

Computes the "level curve" for a function defined as a collection of values on the vertices of a triangular mesh.
# Example

```julia-repl
julia> levelcurve(points,faces,values,level)
2×4 Array{Float64,2}:
 0.0  0.5  0.5  0.0
 0.5  1.0  1.0  0.5

julia> points = [[0 0];[1 0];[0 1];[1 1]]; faces = [[1 2 3];[2 4 3]]; values = [0,1,1,0]; level =1/2;

julia> levelcurve(points,faces,values,level)
2×4 Array{Float64,2}:
 0.0  0.5  0.5  0.0
 1.0  0.5  0.5  1.0
```

In the example above, each row of the matrix correspond to a level edge.
Columns 1 and 2  are the coordinates of the first vertex, and
columns 3 and 4 are the coordinates of the second vertex.
The edges are oriented so that the gradient of the function points to their left side.
"""
function levelcurve(ps::Array{<:Real,2}, fs::Array{<:Int,2}, vals::Array{<:Real,1}, level::Real)
    @assert size(fs,2) == 3
    dim = size(ps,2)
    law =[[[3 1],[3 2]], [[2 3],[2 1]], [[1 3],[1 2]],[[1 2],[1 3]],[[2 1],[2 3]],[[3 2],[3 1]]]
    cnt = 0
    es = zeros(size(fs,1),2*dim) # make enough space to hold edges (edges as pairs of points in 3-space)
    for i in 1:size(fs,1)
        fvs = fs[i,:] # Face vertices
        vals_fvs= vals[fvs] .> level
        ss = 1*vals_fvs[3] + 2*vals_fvs[2] + 4*vals_fvs[1]
        if 0 < ss < 7 # When intersection occurs
            ea, eb = law[ss]
            a1, a2 = fvs[ea]
            b1, b2 = fvs[eb]
            val_a1 = vals[a1]; val_a2 = vals[a2]
            val_b1 = vals[b1]; val_b2 = vals[b2]
            λa = (level-val_a1)/(val_a2-val_a1)
            λb = (level-val_b1)/(val_b2-val_b1)
            pa = (1-λa) *ps[a1,:] + λa*ps[a2,:]
            pb = (1-λb) *ps[b1,:] + λb*ps[b2,:]
            es[cnt+1,1:dim] = pa
            es[cnt+1,dim+1:2*dim] = pb
            cnt = cnt + 1
        end
    end
    es[1:cnt,:]
end

"""
    levelsurface(ps::Array{<:Real,2}, ts::Array{<:Int,2}, vals::Array{<:Real,1}, level::Real)

Computes the "level surface" for a function defined as a collection of values on the vertices of a tetrahedral mesh.
# Example

```julia-repl
julia> points = [[0 0 0];[1 0 0];[0 1 0];[0 0 1]]; tetrahedrons= [[1 2 3 4];]; values = [0,1,1,1]; level =1/2;

julia> levelsurface(points, tetrahedrons, val
valtype values
julia> levelsurface(points, tetrahedrons, values, level)
1×9 Array{Float64,2}:
 0.5  0.0  0.0  0.0  0.5  0.0  0.0  0.0  0.5
```

In the example above, each row of the matrix correspond to a level triangle.
Columns 1, 2 and 3 are the coordinates of the first vertex,
columns 4, 5 and 6 are the coordinates of the second vertex, and
columns 7, 8 and 9 are the coordinates of the third vertex.
The triangles are oriented so that the gradient of the function points to their left side.
"""
function levelsurface(ps::Array{<:Real,2}, ts::Array{<:Int,2}, vals::Array{<:Real,1}, level::Real)
    @assert size(ts,2) == 4
    dim = size(ps,2)
    # vertex1+
    l08 = [[[1 2],[1 4],[1 3]]]
    l07 = [[[1 2],[1 3],[1 4]]]
    # vertex2+
    l04 = [[[2 1],[2 3],[2 4]]]
    l11 = [[[2 1],[2 4],[2 3]]]
    # vertex3+
    l02 = [[[3 1],[3 4],[3 2]]]
    l13 = [[[3 1],[3 2],[3 4]]]
    # vertex4+
    l01 = [[[4 1],[4 2],[4 3]]]
    l14 = [[[4 1],[4 3],[4 2]]]

    #vertices??+
    l12 = [[[1 3],[2 3],[2 4]], [[1 3],[2 4],[1 4]]]
    l03 = [[[1 3],[2 4],[2 3]], [[1 3],[1 4],[2 4]]]
    #vertices??+
    l10 = [[[1 2],[1 4],[4 3]], [[1 2],[4 3],[2 3]]]
    l05 = [[[1 2],[4 3],[1 4]], [[1 2],[2 3],[4 3]]]
    #vertices??+
    l09 = [[[1 2],[2 4],[1 3]], [[1 3],[2 4],[4 3]]]
    l06 = [[[1 2],[1 3],[2 4]], [[1 3],[4 3],[2 4]]]

    ll = [l01,l02,l03,l04,l05,l06,l07,l08,l09,l10,l11,l12,l13,l14]
    cnt = 0
    fs = zeros(size(ts,1),3*dim) # make enough space to hold edges (edges as triplets of points in 3-space)

    for i in 1:size(ts,1)
        tvs = ts[i,:] # Tetrahedron vertices
        vals_tvs= vals[tvs] .> level
        c0, c1, c2, c3 = vals[tvs] .> level
        ss = 1*vals_tvs[4] + 2*vals_tvs[3] + 4*vals_tvs[2] + 8*vals_tvs[1]
        if 0 < ss < 15 # When intersection occurs
            for tr in ll[ss]
                ea, eb, ec = tr # edges having intersection points forming a level face
                a1, a2 = tvs[ea]
                b1, b2 = tvs[eb]
                c1, c2 = tvs[ec]
                val_a1 = vals[a1]; val_a2 = vals[a2]
                val_b1 = vals[b1]; val_b2 = vals[b2]
                val_c1 = vals[c1]; val_c2 = vals[c2]
                λa = (level-val_a1)/(val_a2-val_a1)
                λb = (level-val_b1)/(val_b2-val_b1)
                λc = (level-val_c1)/(val_c2-val_c1)
                pa = (1-λa) *ps[a1,:] + λa*ps[a2,:]
                pb = (1-λb) *ps[b1,:] + λb*ps[b2,:]
                pc = (1-λc) *ps[c1,:] + λc*ps[c2,:]
                fs[cnt+1,1:dim] = pa
                fs[cnt+1,dim+1:2*dim] = pb
                fs[cnt+1,2*dim+1:3*dim] = pc
                cnt = cnt + 1
            end
        end
    end
    fs[1:cnt,:]
end


"""
    abstract type DiffusionParameters

Abstract type for all input structures to diffusion functions.
"""
abstract type DiffusionParameters end

"""
    DiffusionKernelParams

Structure that contains the dependent coefficients used in most of the diffusion theory calculations
"""
struct DiffusionKernelParams{T <: AbstractFloat} <: DiffusionParameters
    D::T                    # Diffusion coefficient
    z0::T                   # isotropic source depth
    ν::T                    # speed of light in medium
    zb::T                   # extrapolation length (EBC)
end
DiffusionKernelParams(μsp::T) where {T<:AbstractFloat} = DiffusionKernelParams{T}(D_coeff(μsp), z0_coeff(μsp), zero(eltype(T)), zero(eltype(T)))
DiffusionKernelParams(μsp::T, n_med::T) where {T<:AbstractFloat} = DiffusionKernelParams{T}(D_coeff(μsp), z0_coeff(μsp), ν_coeff(n_med), zero(eltype(T)))
DiffusionKernelParams(μsp::T, n_med::T, n_ext::T) where {T<:AbstractFloat} = DiffusionKernelParams{T}(D_coeff(μsp), z0_coeff(μsp), ν_coeff(n_med), zb_coeff(A_coeff(n_med / n_ext), D_coeff(μsp)))

"""
    A_coeff(n)

Function to approximate the A coefficient (mismatch between external and internal index of refractions) given in Eq. 25 from Ref. [1].
n is the ratio between the refractive index of the diffusing medium and that of the external medium (n_med / n_ext).
Uses the polynomial fit given in Appendix A [1]. Used to calculate the extrapolation length.

References:
[1] Daniele Contini, Fabrizio Martelli, and Giovanni Zaccanti,
    "Photon migration through a turbid slab described by a model based on diffusion approximation. I. Theory,"
    Appl. Opt. 36, 4587-4599 (1997)
"""
function A_coeff(n)
    if n > 1.0
        A = 504.332889 - 2641.00214n + 5923.699064n^2 - 7376.355814n^3 +
        5507.53041n^4 - 2463.357945n^5 + 610.956547n^6 - 64.8047n^7
    elseif n < 1.0
        A = 3.084635 - 6.531194n + 8.357854n^2 - 5.082751n^3
    else
        A = 1.0
    end
    return A
end

"""
    D_coeff(μsp)

Returns the diffusion coefficient as a function of the reduced scattering coefficient.
Note that this is sometimes defined with the absorption coefficient D = 1 / (3 * (μsp + μa)).
"""
D_coeff(μsp) = 1 / (3 * μsp)

"""
    ν_coeff(n_med)

Computes the speed of light in the medium (cm/ns).
This is a good place to change the units of the returning function along with optical properties if interested in mm or ps.
"""
ν_coeff(n_med) = 29.9792458 / n_med

"""
    z0_coeff(μsp)

The location of the isotropic point source within the medium assuming some perpindicular incidence of light beam.
Note that this is sometimes defined with the absorption coefficient z0 = 1 / (μsp + μa)
"""
z0_coeff(μsp) = 1 / μsp

"""
    zb_coeff(A, D)

The extrapolation length calculated with the A coefficient (accounts for mismatch index of refraction at boundary)
and D (diffusion coefficient).
"""
zb_coeff(A, D) = 2 * A * D


#: Abstract functional types

export MinimizableFunction, DifferentiableFunction, ProximableFunction, ProjectionableSet
export OptimOptions, minimize, minimize!, proxy, proxy!, project, project!, grad, grad!

# Abstract type declarations

"""Expected behavior for MinimizableFunction:
- minimize!(f::MinimizableFunction{T,N}, x0::AbstractArray{T,N}) where {T,N}
It approximates the solution of the optimization problem: min_x f(x)
"""
abstract type MinimizableFunction{T,N} end

abstract type OptimOptions end

minimize(fun::MinimizableFunction{T,N}, x0::AbstractArray{T,N}, opt::OptimOptions) where {T,N} = minimize!(fun, x0, opt, similar(x0))
minimize_debug(fun::MinimizableFunction{T,N}, x0::AbstractArray{T,N}, opt::OptimOptions) where {T,N} = minimize_debug!(fun, x0, opt, similar(x0))

"""Expected behavior for DifferentiableFunction:
- fval::T = grad!(f::DifferentiableFunction{DT}, x::DT, g::DT) where {T,N,DT<:AbstractArray{T,N}}
"""
abstract type DifferentiableFunction{T,N} end

function grad(f::DifferentiableFunction{T,N}, x::AbstractArray{T,N}) where {T,N}
    g = similar(x)
    fval = grad!(f, x, g)
    return fval, g
end

(f::DifferentiableFunction{T,N})(x::AbstractArray{T,N}) where {T,N} = grad!(f, x, nothing)

"""Expected behavior for ProximableFunction:
- proxy!(y::AbstractArray{T,N}, λ::T, g::ProximableFunction{T,N}, x::AbstractArray{T,N}) where {T,N}
- proxy!(y::AbstractArray{T,N}, λ::T, g::ProximableFunction{T,N}, x::AbstractArray{T,N}, opt::OptimOptions) where {T,N}
It approximates the solution of the optimization problem: min_x 0.5*norm(x-y)^2+λ*g(x)
- project!(y::AbstractArray{T,N}, ε::T, g::ProximableFunction{T,N}, x::AbstractArray{T,N}) where {T,N}
- project!(y::AbstractArray{T,N}, ε::T, g::ProximableFunction{T,N}, x::AbstractArray{T,N}, opt::OptimOptions) where {T,N}
It approximates the solution of the optimization problem: min_{g(x)<=ε} 0.5*norm(x-y)^2
"""
abstract type ProximableFunction{T,N} end

proxy(y::AbstractArray{CT,N}, λ::T, g::ProximableFunction{CT,N}) where {T<:Real,N,CT<:RealOrComplex{T}} = proxy!(y, λ, g, similar(y))
project(y::AbstractArray{CT,N}, ε::T, g::ProximableFunction{CT,N}) where {T<:Real,N,CT<:RealOrComplex{T}} = project!(y, ε, g, similar(y))
proxy(y::AbstractArray{CT,N}, λ::T, g::ProximableFunction{CT,N}, opt::OptimOptions) where {T<:Real,N,CT<:RealOrComplex{T}} = proxy!(y, λ, g, similar(y), opt::OptimOptions)
project(y::AbstractArray{CT,N}, ε::T, g::ProximableFunction{CT,N}, opt::OptimOptions) where {T<:Real,N,CT<:RealOrComplex{T}} = project!(y, ε, g, similar(y), opt::OptimOptions)

"""Projectional sets
Expected behavior for convex sets: y = project!(x, C, y), y = Π_C(x)
"""
abstract type ProjectionableSet{T,N} end

project(x::AbstractArray{T,N}, C::ProjectionableSet{T,N}) where {T,N} = project!(x, C, similar(x))
project(x::AbstractArray{T,N}, C::ProjectionableSet{T,N}, opt::OptimOptions) where {T,N} = project!(x, C, similar(x), opt)


export stablepartition!, isstableset

function isstableset(S::Set{T}, E::Set{Tuple{T,T}}) where T
  for (i,j) in E
    if i in S && j in S
      return false
    end
  end
  return true
end

function stablepartition!(conflicts::Set{Tuple{T, T}},
                             stablesets::Array{Set{T}, 1} = Set{T}[])  where T

  # elements already in stablesets are already taken care of
  prepartitioned = union(stablesets...)

  # convert the set of conflicts into a dictionary
  # mapping each id to the ids of its neighbors in the graph
  id2neighborid = Dict{T, Set{T}}() # map from variable id to neighbor ids
  for e in conflicts
    for i in 1:2
      v = e[i]
      n = e[i%2+1]
      if v in prepartitioned
        continue
      elseif !(v in keys(id2neighborid))
        # initialize v and add neighbor n
        id2neighborid[v] = Set{T}([n])
      else
        # just add neighbor n
        union!(id2neighborid[v], n)
      end
    end
  end

  # apportion variables not prepartitioned into the stable sets,
  # creating additional stable sets if necessary
  # might be better to produce a set cover with maximal stable sets,
  # but that's (NP) hard
  stablepartition!(id2neighborid, stablesets)

  return stablesets
end

# id2neighborhoodid maps node ids to the list of the node's neighbors in the graph
# returns a list of stable sets in the graph whose union is the full set of nodes
# uses a greedy algorithm
function stablepartition!(id2neighborid::Dict{T, Set{T}},
                             stablesets::Array{Set{T}, 1} = Set{T}[])  where T
  for varid in keys(id2neighborid)
    foundhome = false
    varid in id2neighborid[varid] && error("cannot partition graph with self loop")
    for s in stablesets
      if length(intersect(s, id2neighborid[varid])) == 0
        # var has no neighbors in s, so add it to s
        union!(s, varid)
        foundhome = true
        break
      end
    end
    if !foundhome
      # var conflicts with all previous subsets, so make a new subset
      push!(stablesets, Set{T}([varid]))
    end
  end
  return stablesets
end


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 97e18b69-4eb6-42a9-beef-04697156a7ab
using PlutoUI, JuMP

# ╔═╡ 2f845732-c095-11eb-14e8-57175fb5038e
lines = readlines(open("./benchmark_problems/Class B/INST13.txt"))

# ╔═╡ 12f3987f-ea86-46a2-852f-65ce5886cbe2
split_line(line) = filter(x->x!="", split(line, " "))

# ╔═╡ 2c521bfc-683b-4337-b62e-a5586d3d06cb
function parse_int(i)
	try
		parse(Int, i)
	catch ArgumentError
		parse(Int, replace(i, ".00"=>""))
	end
end

# ╔═╡ d4175f5d-d4e9-4862-9c60-83a74af1dc1c
parse_line(line) = map(n->parse_int(n), split_line(line))

# ╔═╡ 0707ad63-038b-4b0f-9c48-310502701915
file = map(parse_line, lines[2:end])

# ╔═╡ dca8d05b-f878-4fb6-bb8e-cef62dd2249a
m_blocks, n_rows, n_groups = file[1]

# ╔═╡ f1d972de-aa4a-4e5e-91d6-890bd4b21fcd
groups_bounds_row = file[2]

# ╔═╡ 75894d1c-b1c4-4050-8b52-19d837e2bc31
data = file[3:end]

# ╔═╡ cf5f3e19-7087-4e67-bb3e-d02e2b33efeb
block_start_index(i) = (i-1)*(n_rows+1) + 1

# ╔═╡ dc9221b9-09b8-48e9-b46b-8cad6cb0b778
@bind row Slider(1:n_groups)

# ╔═╡ 39ba33bb-7113-4f35-9f19-676178e210cb
n_groups

# ╔═╡ c9fe0f02-8f6a-4584-88ca-d0640fab9be2
n_i = n_rows

# ╔═╡ 2ef7e159-9bff-4b36-92a3-52c43db1c8fe
m = n_groups

# ╔═╡ d7205749-58ab-4831-8797-ac3d7a175f37
falses(m)

# ╔═╡ c06e268b-09a2-482f-98fa-6ecbd8ee9fe1
n = m_blocks # size of set of disjoint groups

# ╔═╡ 9c4c419c-0517-4ae9-9a48-7368444085fe
function load_block(i, n_groups, data)
	d = data[block_start_index(i):block_start_index(i+1)-1]
	@assert first(first(d)) == i
	items = d[2:end]
	#exactly one of each element in this group must be turned on
	#in order to represent this in MPS form, we need to one-hot encode the class index
	index = falses(n_groups)
	index[i] = true

	#then combine the one-hot encoding with the resource requirements
	resource_reqs = map(x->vcat(x[2:end], index), items)

	#strip the profit vector, then return the (profit, resources) tuple
	profit = map(x->first(x), items)
	profit, resource_reqs
end

# ╔═╡ 1c10debb-ee04-45e2-81fc-9fe4d9362bc7
current_block = load_block(row, n_groups, data)

# ╔═╡ e764cd4e-fc74-4561-8265-e303b2892c88
length(current_block)

# ╔═╡ 8652e037-b90b-4044-a618-730f1f554b2e
function make_model(profit, resource_reqs, available_res)
	m = Model()
	x = @variable(m, x[1:length(profit)])
	add_resource_constraint(m, x, r, a) =
		@constraint(m, sum(x .* r) <= a)
	map((i, row)->add_resource_constraint(m, x, row, a[i]), resource_reqs)
	@objective(m, Max, sum(x .* profit))
	m
end

# ╔═╡ d3938ef5-e05a-47c3-b4c6-bc966245acd8
blocks = map(i->load_block(i, n_groups, data), 1:n_groups)

# ╔═╡ c35b78db-1c40-4313-a580-83f48680a7e1
profits = vcat(map(b->first(b), blocks)...)

# ╔═╡ 1db24586-4b11-4b75-bfee-1015644527de
resource_reqs = vcat(map(b->b[2], blocks)...)

# ╔═╡ 8fd748aa-6aff-4dcc-a5fd-c50d52295dc5
resource_limits = vcat(groups_bounds_row..., ones(n_groups))

# ╔═╡ Cell order:
# ╠═2f845732-c095-11eb-14e8-57175fb5038e
# ╠═12f3987f-ea86-46a2-852f-65ce5886cbe2
# ╠═d4175f5d-d4e9-4862-9c60-83a74af1dc1c
# ╠═2c521bfc-683b-4337-b62e-a5586d3d06cb
# ╠═0707ad63-038b-4b0f-9c48-310502701915
# ╠═dca8d05b-f878-4fb6-bb8e-cef62dd2249a
# ╠═f1d972de-aa4a-4e5e-91d6-890bd4b21fcd
# ╠═75894d1c-b1c4-4050-8b52-19d837e2bc31
# ╠═cf5f3e19-7087-4e67-bb3e-d02e2b33efeb
# ╠═1c10debb-ee04-45e2-81fc-9fe4d9362bc7
# ╠═d7205749-58ab-4831-8797-ac3d7a175f37
# ╠═dc9221b9-09b8-48e9-b46b-8cad6cb0b778
# ╠═39ba33bb-7113-4f35-9f19-676178e210cb
# ╠═e764cd4e-fc74-4561-8265-e303b2892c88
# ╠═c9fe0f02-8f6a-4584-88ca-d0640fab9be2
# ╠═2ef7e159-9bff-4b36-92a3-52c43db1c8fe
# ╠═c06e268b-09a2-482f-98fa-6ecbd8ee9fe1
# ╠═9c4c419c-0517-4ae9-9a48-7368444085fe
# ╠═8652e037-b90b-4044-a618-730f1f554b2e
# ╠═d3938ef5-e05a-47c3-b4c6-bc966245acd8
# ╠═c35b78db-1c40-4313-a580-83f48680a7e1
# ╠═1db24586-4b11-4b75-bfee-1015644527de
# ╠═8fd748aa-6aff-4dcc-a5fd-c50d52295dc5
# ╠═97e18b69-4eb6-42a9-beef-04697156a7ab