text:
# ----------------------------------------------------------------------------------------------------------
# Model generation
# ----------------------------------------------------------------------------------------------------------
function model_DRO(n::Int, d::Projections.Constraint, ε::Real = 0.1; seed::Int = 1234)
    Random.seed!(seed);

    q = LinearAlgebra.normalize(rand(Distributions.Uniform(0,1), n), 1)
    c = rand(Distributions.Normal(0,1), n)
    return Projections.DRO(d, q, c, ε)
end

function model_Simplex(n::Int; seed::Int = 1234)
    Random.seed!(seed);

    q  = rand(n)
    lb = rand(n)./n
    ub = 1 .+ rand(n)
    return Projections.Simplex(q, lb, ub)
end

# ----------------------------------------------------------------------------------------------------------
# Benchmark evaluation
# ----------------------------------------------------------------------------------------------------------
function eval_Divergences(solver::Solver, N; kwargs...)
    divergences = [KullbackLeibler(), Hellinger(), Burg(), ChiSquare(), ModifiedChiSquare()]

    return  mapreduce((d) -> eval_DRO(solver, d, N; kwargs...), vcat, divergences)
end

function eval_Norms(solver::Solver, N; kwargs...)
    norms = [Linf(), Lone(), Ltwo()]

    return  mapreduce((d) -> eval_DRO(solver, d, N; kwargs...), vcat, norms)
end

function eval_DRO(solver::Solver, d, N; kwargs...)
    return  Projections.benchmark(solver, (n) -> model_DRO(n, d), N; kwargs...)
end

function eval_Simplex(solver::Solver, N; kwargs...)
    return Projections.benchmark(solver, model_Simplex, N; kwargs...)
end

# ----------------------------------------------------------------------------------------------------------
# Tables
# ----------------------------------------------------------------------------------------------------------
function tableformat(metric::Symbol)
    colnames = [:dimension, :dimension_small]
    for problem in [:KullbackLeibler, :Hellinger, :Burg, :ChiSquare, :ModifiedChiSquare]
        push!(colnames, Symbol(problem, "_Our"))
        metric in [:evals_mean, :evals_std] || push!(colnames, Symbol(problem, "_General"))
    end

    if metric in [:evals_mean, :evals_std]
        push!(colnames, :Ltwo_Our)
    else
        for problem in [:Linf, :Lone, :Ltwo, :Simplex]
            push!(colnames, Symbol(problem, "_Our"))
            push!(colnames, Symbol(problem, "_General"))
            problem == :Ltwo && push!(colnames, Symbol(problem, "_Philpott"))
        end
    end
    return colnames
end

function savetable(table::DataFrame;
                   save::Bool       = false,
                   savepath::String = "",
                   savename::String = "comparison")

    savetable(table, :evaltime_mean; save = save, savepath = savepath, savename = savename)
    savetable(table, :evaltime_std;  save = save, savepath = savepath, savename = savename)
    savetable(table, :evals_mean;    save = save, savepath = savepath, savename = savename)
    savetable(table, :evals_std;     save = save, savepath = savepath, savename = savename)

    save && CSV.write(joinpath(savepath,  "$(savename)_full_dataframe.csv"), table)

    return
end

function savetable(table::DataFrame,
                   metric::Symbol;
                   save::Bool       = false,
                   savepath::String = "",
                   savename::String = "comparison")

    cols      = [:model, :constraint, :solver]

    dimension = sort(unique(table.dimension))
    n         = length(dimension)
    table_new = DataFrames.DataFrame(dimension = dimension)

    dimension_small = sort(unique(table.dimension[table.solver .== "General"]))
    l               = length(dimension_small)
    table_new[!, :dimension_small] = vcat(dimension_small, [missing for i in 1:n-l])

    for df in DataFrames.groupby(table, cols)
        colname  = join([string(df[1, col]) for col in cols], "_")
        colvalue = df[!, metric]
        l        = length(colvalue)

        table_new[!, Symbol("$(colname)")] = vcat(colvalue, [missing for i in 1:n-l])
    end

    colnames = [replace(colname, "DRO_" => "")  for colname in string.(names(table_new))]
    colnames = [replace(colname, "none_" => "") for colname in colnames]
    names!(table_new, Symbol.(colnames))

    table_new = table_new[!, tableformat(metric)]

    save && CSV.write(joinpath(savepath,  "$(savename)_$(metric).csv"), table_new)
    return table_new
end

# ----------------------------------------------------------------------------------------------------------
# Plots
# ----------------------------------------------------------------------------------------------------------
function comparison(table::DataFrames.DataFrame,
                    varname::Symbol  = :evaltime;
                    xscale::Symbol   = :identity,
                    yscale::Symbol   = :identity,
                    ploterror::Bool  = false,
                    title::String    = "$(varname) comparison",
                    save::Bool       = false,
                    savename::String = title,
                    savepath::String = "",
                    type::String     = "png")

    cols = [:model, :constraint, :solver, :optimizer]

    p = plot(size = (900, 600), dpi = 150, legend = :topleft)
    title!(title)
    xlabel!("n")
    ylabel!(string(varname))
    for df in DataFrames.groupby(table, cols)
        label = join([df[1, col] for col in cols], " - ")

        n     = df.dimension
        t     = df[!, Symbol(varname, :_mean)]
        t_std = df[!, Symbol(varname, :_std)]

        if ploterror
            plot!(n, t, yerror = t_std, label = label, xscale = xscale, yscale = yscale)
        else
            plot!(n, t, label = label, xscale = xscale, yscale = yscale)
        end
    end

    save && savefig(p, joinpath(savepath, "$(savename).$(type)"))
    return p
end

function h(λ::AbstractVector{<:Real},
           m::T;
           title::String    = "h_$(typeof(m.d).name)",
           save::Bool       = false,
           savename::String = title,
           savepath::String = "",
           type::String     = "png") where {T <: Projections.Model}

    f(λ) = Projections.h(m.d, m, λ)

    p = plot(size = (900, 600), dpi = 150, legend = :best, show = false)
    title!(title)
    xlabel!("lambda")
    ylabel!("h(lambda)")
    plot!(λ, f.(λ))

    save && savefig(p, joinpath(savepath, "$(savename).$(type)"))
    return p
end

# ----------------------------------------------------------------------------------------------------------
# Comparison solver
# ----------------------------------------------------------------------------------------------------------
function comparison_solver(solver::T,
                           N::AbstractVector{<:Int};
                           save::Bool        = false,
                           savepath::String  = "",
                           maxevals::Integer = 10) where {T<:Solver}

    ## DRO with divergences
    @info "DRO with divergences"
    table1 = eval_Divergences(solver, N; verbose = true,  maxevals = maxevals)

    ## DRO with norms
    @info "DRO with norms"
    table2 = eval_Norms(solver, N; verbose = true,  maxevals = maxevals)

    ## Simplex1
    @info "Simplex model"
    table3 = eval_Simplex(solver, N; verbose = true,  maxevals = maxevals)

    ## Plots - time comparison
    comparison(table2, :evaltime; save = save, savepath = savepath, savename = "time_norms_$(T.name)")
    comparison(table1, :evaltime; save = save, savepath = savepath, savename = "time_divergences_$(T.name)")
    comparison(table3, :evaltime; save = save, savepath = savepath, savename = "time_simplex_$(T.name)")

    ## Plots - #objective function evaluation
    if T <: Our
        comparison(table2, :evals;    save = save, savepath = savepath, savename = "evals_norms_$(T.name)")
        comparison(table1, :evals;    save = save, savepath = savepath, savename = "evals_divergences_$(T.name)")
        comparison(table3, :evals;    save = save, savepath = savepath, savename = "evals_simplex_$(T.name)")
    end

    return vcat(table1, table2, table3)
end

# ----------------------------------------------------------------------------------------------------------
# Comparison Philpott
# ----------------------------------------------------------------------------------------------------------
function comparison_philpott(N::AbstractVector{<:Int};
                             save::Bool        = false,
                             savepath::String  = "",
                             maxevals::Integer = 10)

    ## DRO with l2 norm and Philpott() solver
    @info "DRO with l2 norm and Philpott() solver"
    table = eval_DRO(Philpott(), Ltwo(), N; verbose = true,  maxevals = maxevals)

    ## Plots - time comparison
    comparison(table, :evaltime; save = save, savepath = savepath, savename = "time_philpott")
    return table
end

max_stars_repo_name: VaclavMacha/Projections

max_issues_repo_name: VaclavMacha/Projections

max_forks_repo_name: VaclavMacha/Projections


text:
"""
    asynchronous(A::AbstractMatrix)

Calculate a asynchronous two-dimensional correlation spectrum.

`A` is a dataset with multiple spectra in rows.
They *need* to be at equally spaced measurements of the perturbation.
Normally, the data is given as the difference from a reference spectrum (see
below).

# Examples
```jldoctest
julia> A = rand(50, 1000);

julia> Ψ = asynchronous(A .- mean(A, dims=1));

julia> size(Ψ)
(1000, 1000)
```
"""
function asynchronous(A::AbstractMatrix)
    return A' * (N * A) / (size(A, 1) - 1)
end

ext: jl


text:
mutable struct ScalarField
    name::String
    type::String
    #time::String
    cellValues::Vector{Float64}
    faceValues::Vector{Float64}
end

mutable struct VectorField
    name::String
    type::String
    #time::String
    cellValues::Vector{Vector{Float64}}
    faceValues::Vector{Vector{Float64}}
end

function setupScalarField(name, value=0.0)
    cellValues = [value for _ in 1:(mesh.nCells + mesh.nBoundaryFaces)]
    faceValues = [value for _ in 1:mesh.nFaces]
    return ScalarField(name, "scalar", cellValues, faceValues)
end

function setupVectorField(name, value=[0.0, 0.0, 0.0])
    cellValues = [value for _ in 1:(mesh.nCells + mesh.nBoundaryFaces)]
    faceValues = [value for _ in 1:mesh.nFaces]
    return VectorField(name, "vector", cellValues, faceValues)
end


text:
"""
    abstract type AbstractLGC end

Represents a Gaussian conditional distribution:
```julia
y | x ∼ Gaussian(A * x + a, Q)
```

Subtypes have discretion over how to implement the interface for this type. In particular
`A`, `a`, and `Q` may not be represented explicitly so that structure can be exploited to
accelerate computations.

# Interface:
- `==`
- `eltype`
- `dim_out`
- `dim_in`
- `predict`
- `predict_marginals`
- `conditional_rand`
- `ε_randn`
- `posterior_and_lml`
"""
abstract type AbstractLGC end

Base.:(==)(x::AbstractLGC, y::AbstractLGC) = (x.A == y.A) && (x.a == y.a) && (x.Q == y.Q)

Base.eltype(f::AbstractLGC) = eltype(f.A)

"""
    predict(x::Gaussian, f::AbstractLGC)

Compute the distribution "predicted" by this conditional given a `Gaussian` input `x`. Will
be equivalent to
```julia
    Gaussian(f.A * x.m + f.a, f.A * x.P * f.A' + f.Q)
```
"""
function predict(x::Gaussian, f::AbstractLGC)
    A, a, Q = get_fields(f)
    m, P = get_fields(x)
    # Symmetric wrapper needed for numerical stability. Do not unwrap.
    return Gaussian(A * m + a, A * symmetric(P) * A' + Q)
end

"""
    predict_marginals(x::Gaussian, f::AbstractLGC)

Equivalent to
```julia
    y = predict(x, f)
    Gaussian(mean(y), Diagonal(cov(y)))
```
"""
function predict_marginals(x::Gaussian, f::AbstractLGC)
    return Gaussian(
        f.A * x.m + f.a,
        Diagonal(diag_At_B(f.A', x.P * f.A') + diag(f.Q)),
    )
end

"""
    conditional_rand(rng::AbstractRNG, f::AbstractLGC, x::AbstractVector)
    conditional_rand(ε::AbstractVector, f::AbstractLGC, x::AbstractVector)

Sample from the conditional distribution `y | x`. `ε` is the randomness needed to generate
this sample. If `rng` is provided, it will be used to construct `ε` via `ε_randn`.

If implementing a new `AbstractLGC`, implement the `ε` method as it avoids randomness, which
means that it plays nicely with `scan_emit`'s checkpointed pullback.
"""
function conditional_rand(rng::AbstractRNG, f::AbstractLGC, x::AbstractVector)
    return conditional_rand(ε_randn(rng, f), f, x)
end

function conditional_rand(ε::AbstractVector, f::AbstractLGC, x::AbstractVector)
    A, a, Q = get_fields(f)
    return (A * x + a) + cholesky(symmetric(Q + UniformScaling(1e-9))).U' * ε
end

"""
    ε_randn(rng::AbstractRNG, f::AbstractLGC)

Generate the vector of random numbers needed inside `conditional_rand`.
"""
ε_randn(rng::AbstractRNG, f::AbstractLGC) = ε_randn(rng, f.A)
ε_randn(rng::AbstractRNG, A::AbstractMatrix{T}) where {T<:Real} = randn(rng, T, size(A, 1))
function ε_randn(rng::AbstractRNG, A::SMatrix{Dout, Din, T}) where {Dout, Din, T<:Real}
    return randn(rng, SVector{Dout, T})
end

Zygote._pullback(::AContext, ::typeof(ε_randn), args...) = ε_randn(args...), nograd_pullback

scalar_type(x::AbstractVector{T}) where {T} = T
scalar_type(x::T) where {T<:Real} = T

Zygote._pullback(::AContext, ::typeof(scalar_type), x) = scalar_type(x), nograd_pullback

"""
    SmallOutputLGC{
        TA<:AbstractMatrix, Ta<:AbstractVector, TQ<:AbstractMatrix,
    } <: AbstractLGC

a.k.a. LGC. An `AbstractLGC` designed for problems in which `A` is a matrix, and
`size(A, 1) < size(A, 2)`. It should still work (roughly) for problems in which
`size(A, 1) > size(A, 2)`, but one should expect worse accuracy and performance than a
`LargeOutputLGC` in such circumstances.
"""
struct SmallOutputLGC{
    TA<:AbstractMatrix, Ta<:AbstractVector, TQ<:AbstractMatrix,
} <: AbstractLGC
    A::TA
    a::Ta
    Q::TQ
end

dim_out(f::SmallOutputLGC) = size(f.A, 1)

dim_in(f::SmallOutputLGC) = size(f.A, 2)

noise_cov(f::SmallOutputLGC) = Zygote.literal_getfield(f, Val(:Q))

function get_fields(f::SmallOutputLGC)
    A = Zygote.literal_getfield(f, Val(:A))
    a = Zygote.literal_getfield(f, Val(:a))
    Q = Zygote.literal_getfield(f, Val(:Q))
    return A, a, Q
end

function posterior_and_lml(x::Gaussian, f::SmallOutputLGC, y::AbstractVector{<:Real})
    m, P = get_fields(x)
    A, a, Q = get_fields(f)

    V = A * P

    S = cholesky(symmetric(V * A' + Q))
    B = S.U' \ V
    α = S.U' \ (y - (A * m + a))

    lml = -(length(y) * convert(scalar_type(y), log(2π)) + logdet(S) + α'α) / 2
    return Gaussian(m + B'α, P - B'B), lml
end

function posterior_and_lml(
    x::Gaussian, f::SmallOutputLGC, y::AbstractVector{<:Union{Missing, <:Real}},
)
    # This implicitly assumes that Q is Diagonal. MethodError if not.
    A, a, Q = get_fields(f)
    Q_filled, y_filled = fill_in_missings(Q, y)
    x_post, lml_raw = posterior_and_lml(x, SmallOutputLGC(A, a, Q_filled), y_filled)
    return x_post, lml_raw + _logpdf_volume_compensation(y)
end

# Required for type-stability. This is a technical detail.
function Zygote._pullback(::NoContext, ::Type{<:SmallOutputLGC}, A, a, Q)
    SmallOutputLGC_pullback(::Nothing) = nothing
    SmallOutputLGC_pullback(Δ) = nothing, Δ.A, Δ.a, Δ.Q
    return SmallOutputLGC(A, a, Q), SmallOutputLGC_pullback
end

# Required for type-stability. This is a technical detail.
function Zygote._pullback(::NoContext, ::typeof(+), A::Matrix{<:Real}, D::Diagonal{<:Real})
    plus_pullback(Δ::Nothing) = nothing
    plus_pullback(Δ) = (nothing, Δ, (diag=diag(Δ),))
    return A + D, plus_pullback
end

"""
    LargeOutputLGC{
        TA<:AbstractMatrix, Ta<:AbstractVector, TQ<:AbstractMatrix,
    } <: AbstractLGC

A SmallOutputLGC (LGC) specialised for models in which the dimension of the
outputs are greater than that of the inputs. These specialisations both improve numerical
stability and performance (time and memory), so it's worth using if your model lives in
this regime.
"""
struct LargeOutputLGC{
    TA<:AbstractMatrix, Ta<:AbstractVector, TQ<:AbstractMatrix,
} <: AbstractLGC
    A::TA
    a::Ta
    Q::TQ
end

function Zygote._pullback(
    ::AContext,
    ::Type{<:LargeOutputLGC},
    A::AbstractMatrix,
    a::AbstractVector,
    Q::AbstractMatrix,
)
    LargeOutputLGC_pullback(Δ) = nothing, Δ.A, Δ.a, Δ.Q
    return LargeOutputLGC(A, a, Q), LargeOutputLGC_pullback
end

dim_out(f::LargeOutputLGC) = size(f.A, 1)

dim_in(f::LargeOutputLGC) = size(f.A, 2)

noise_cov(f::LargeOutputLGC) = Zygote.literal_getfield(f, Val(:Q))

function get_fields(f::LargeOutputLGC)
    A = Zygote.literal_getfield(f, Val(:A))
    a = Zygote.literal_getfield(f, Val(:a))
    Q = Zygote.literal_getfield(f, Val(:Q))
    return A, a, Q
end

function posterior_and_lml(x::Gaussian, f::LargeOutputLGC, y::AbstractVector{<:Real})
    m, _P = get_fields(x)
    A, a, _Q = get_fields(f)
    Q = cholesky(symmetric(_Q))
    P = cholesky(symmetric(_P + ident_eps(1e-10)))

    # Compute posterior covariance matrix.
    Bt = Q.U' \ A * P.U'
    F = cholesky(symmetric(Bt' * Bt + UniformScaling(1.0)))
    G = F.U' \ P.U
    P_post = G'G

    # Compute posterior mean.
    δ = Q.U' \ (y - (A * m + a))
    β = F.U' \ (Bt' * δ)
    m_post = m + G' * β

    # Compute log marginal likelihood.
    c = convert(scalar_type(y), length(y) * log(2π))
    lml = _compute_lml(δ, F, β, c, Q)

    return Gaussian(m_post, P_post), lml
end

# For some compiler-y reason, chopping this up helps.
_compute_lml(δ, F, β, c, Q) = -(δ'δ - β'β + c + logdet(F) + logdet(Q)) / 2

function posterior_and_lml(
    x::Gaussian, f::LargeOutputLGC, y::AbstractVector{<:Union{Missing, <:Real}},
)
    A, a, Q = get_fields(f)
    # This implicitly assumes that Q is Diagonal. MethodError if not.
    Q_filled, y_filled = fill_in_missings(Q, y)
    x_post, lml_raw = posterior_and_lml(x, LargeOutputLGC(A, a, Q_filled), y_filled)
    return x_post, lml_raw + _logpdf_volume_compensation(y)
end

"""
    ScalarOutputLGC

An LGC that operates on a vector-valued input space and a scalar-valued output space.
Similar to `SmallOutputLGC` when its `dim_out` is 1 but, for example, `conditional_rand`
returns a `Real` rather than an `AbstractVector` of length 1.
"""
struct ScalarOutputLGC{
    TA<:Adjoint{<:Any, <:AbstractVector}, Ta<:Real, TQ<:Real,
} <: AbstractLGC
    A::TA
    a::Ta
    Q::TQ
end

dim_out(f::ScalarOutputLGC) = 1

dim_in(f::ScalarOutputLGC) = size(f.A, 2)

function get_fields(f::ScalarOutputLGC)
    A = Zygote.literal_getfield(f, Val(:A))
    a = Zygote.literal_getfield(f, Val(:a))
    Q = Zygote.literal_getfield(f, Val(:Q))
    return A, a, Q
end

noise_cov(f::ScalarOutputLGC) = Zygote.literal_getfield(f, Val(:Q))

function conditional_rand(ε::Real, f::ScalarOutputLGC, x::AbstractVector)
    A, a, Q = get_fields(f)
    return (A * x + a) + sqrt(Q) * ε
end

ε_randn(rng::AbstractRNG, f::ScalarOutputLGC) = randn(rng, eltype(f))

function posterior_and_lml(x::Gaussian, f::ScalarOutputLGC, y::T) where {T<:Real}
    m, P = get_fields(x)
    A, a, Q = get_fields(f)
    V = A * P
    sqrtS = sqrt(V * A' + Q)
    B = sqrtS \ V
    α = sqrtS \ (y - (A * m + a))

    lml = -(convert(T, log(2π)) + 2 * log(sqrtS) + α^2) / 2
    return Gaussian(m + B'α, P - B'B), lml
end

"""
    BottleneckLGC

A composition of an affine map that projects onto a low-dimensional subspace and a
`LargeOutputLGC`. This structure is exploited by only ever computing `Cholesky`
factorisations in the space the affine map maps to, rather than the input or output space.

Letting, `H` and `h` parametrise the affine map, and `f` the "fan-out" `LargeOutputLGC`, the
conditional distribution that this model parametrises is
```julia
y | x ~ Gaussian(f.A * (H * x + h) + f.a, f.Q)
```

Note that this type does not enforce that `size(H, 1) < size(H, 2)`, nor that
`dim_out(f) > dim_in(f)`, it's just not a particularly good idea to use this type unless
your model satisfies these properties.
"""
struct BottleneckLGC{
    TH<:AbstractMatrix{<:Real}, Th<:AbstractVector{<:Real}, Tout<:LargeOutputLGC,
} <: AbstractLGC
    H::TH
    h::Th
    fan_out::Tout
end

function Base.:(==)(x::BottleneckLGC, y::BottleneckLGC)
    return (x.H == y.H) && (x.h == y.h) && (x.fan_out == y.fan_out)
end

Base.eltype(f::BottleneckLGC) = eltype(f.fan_out)

dim_out(f::BottleneckLGC) = dim_out(f.fan_out)

dim_in(f::BottleneckLGC) = size(f.H, 2)

noise_cov(f::BottleneckLGC) = noise_cov(Zygote.literal_getfield(f, Val(:fan_out)))

function get_fields(f::BottleneckLGC)
    H = Zygote.literal_getfield(f, Val(:H))
    h = Zygote.literal_getfield(f, Val(:h))
    fan_out = Zygote.literal_getfield(f, Val(:fan_out))
    return H, h, fan_out
end

fan_out(f::BottleneckLGC) = Zygote.literal_getfield(f, Val(:fan_out))

function conditional_rand(ε::AbstractVector{<:Real}, f::BottleneckLGC, x::AbstractVector)
    H, h, fan_out = get_fields(f)
    return conditional_rand(ε, fan_out, H * x + h)
end

ε_randn(rng::AbstractRNG, f::BottleneckLGC) = ε_randn(rng, f.fan_out)

# Construct the low-dimensional projection given by f.H and f.h of `x`.
function _project(x::Gaussian, f::BottleneckLGC)
    m, P = get_fields(x)
    H, h, _ = get_fields(f)
    return Gaussian(H * m + h, H * P * H' + ident_eps(x))
end

function predict(x::Gaussian, f::BottleneckLGC)
    return predict(_project(x, f), fan_out(f))
end

function predict_marginals(x::Gaussian, f::BottleneckLGC)
    return predict_marginals(_project(x, f), fan_out(f))
end

function posterior_and_lml(x::Gaussian, f::BottleneckLGC, y::AbstractVector)

    xm, xP = get_fields(x)
    H, _, fan_out = get_fields(f)

    # Get the posterior over the intermediate variable `z`.
    z = _project(x, f)
    z_post, lml = posterior_and_lml(z, fan_out, y)

    # Compute the posterior `x | y` by integrating `x | z` against `z | y`.
    zm, zP = get_fields(z)
    z_postm, z_postP = get_fields(z_post)
    U = cholesky(symmetric(zP + ident_eps(z, 1e-12))).U
    Gt = U \ (U' \ (H * xP))
    return Gaussian(xm + Gt' * (z_postm - zm), xP + Gt' * (z_postP - zP) * Gt), lml
end


text:
using Test

#Arithmetic library imports
using IntervalArithmetic
using IntervalContractors

#Preamble
setprecision(53)
setprecision(Interval, Float64)

# Using approximate checks for validaty update later?
import Base.isapprox
isapprox(x::Interval,y::Interval) = isapprox(x.lo,y.lo,atol=1E-4) && isapprox(x.hi,y.hi,atol=1E-4)

@testset "max_rev_test" begin
    @test max_rev(∅, entireinterval(Float64),entireinterval(Float64))[2] == Interval(-∞, ∞)
    @test max_rev(Interval(0.0, 1.0), ∅,Interval(-2.0, -1.0))[2] == ∅
    @test max_rev(∅, Interval(0.0, 1.0),Interval(-2.0, -1.0))[2] == ∅
    @test max_rev(Interval(-2.0, -1.0), entireinterval(Float64),Interval(-3.0, -2.0))[2] == Interval(-∞, -1.0)
    @test isapprox(max_rev(Interval(1.0, 1.0), entireinterval(Float64),Interval(-2.0, -1.0))[2], Interval(-∞, 1))
    @test max_rev(entireinterval(Float64), Interval(-2.0, -1.0), Interval(-3.0, -2.0))[2] == Interval(-2.0, -1.0)
    @test max_rev(Interval(-Inf, 0.0), entireinterval(Float64),Interval(-3.0, -2.0))[2] == Interval(-∞, 0.0)
end

@testset "min_rev_test" begin
    @test min_rev(∅, entireinterval(Float64),entireinterval(Float64))[2] == Interval(-∞, ∞)
    @test min_rev(Interval(0.0, 1.0), ∅,Interval(-2.0, -1.0))[2] == ∅                             # should return empty?
    @test min_rev(∅, Interval(0.0, 1.0),Interval(-2.0, -1.0))[2] == ∅
    @test min_rev(Interval(-2.0, -1.0), entireinterval(Float64),Interval(-3.0, -2.0))[2] == Interval(-∞, -2.0)
    @test isapprox(min_rev(Interval(1.0, 1.0), entireinterval(Float64),Interval(-2.0, -1.0))[2], Interval(-∞, -1.0))
    @test min_rev(entireinterval(Float64), Interval(-2.0, -1.0), Interval(-3.0, -2.0))[2] == Interval(-3.0, -2.0)
    @test min_rev(Interval(-Inf, 0.0), entireinterval(Float64),Interval(-3.0, -2.0))[2] == Interval(-∞, -2.0)
end


text:
# Autogenerated wrapper script for faust_jll for x86_64-linux-gnu-cxx11
export faust, faust2alqt, faust2alsa, faust2alsaconsole, faust2android, faust2androidunity, faust2api, faust2atomsnippets, faust2au, faust2bela, faust2caqt, faust2caqtios, faust2csound, faust2csvplot, faust2dplug, faust2dssi, faust2dummy, faust2dummymem, faust2eps, faust2esp32, faust2faustvst, faust2firefox, faust2gen, faust2graph, faust2graphviewer, faust2ios, faust2jack, faust2jackconsole, faust2jackrust, faust2jackserver, faust2jaqt, faust2jaqtchain, faust2juce, faust2ladspa, faust2linuxunity, faust2lv2, faust2mathdoc, faust2mathviewer, faust2max6, faust2md, faust2msp, faust2netjackconsole, faust2netjackqt, faust2nodejs, faust2octave, faust2osxiosunity, faust2owl, faust2paqt, faust2pdf, faust2plot, faust2png, faust2portaudiorust, faust2pure, faust2puredata, faust2raqt, faust2ros, faust2rosgtk, faust2rpialsaconsole, faust2rpinetjackconsole, faust2sam, faust2sc, faust2sig, faust2sigviewer, faust2smartkeyb, faust2sndfile, faust2soul, faust2supercollider, faust2svg, faust2svgviewer, faust2teensy, faust2unity, faust2vcvrack, faust2vst, faust2vsti, faust2w32max6, faust2w32msp, faust2w32puredata, faust2w32vst, faust2w64max6, faust2w64vst, faust2wasm, faust2webaudiowasm, faust2webaudiowast, faust2winunity, faustoptflags, libHTTPDFaust, libOSCFaust, libfaust, stdfaust_lib, usage_sh

using LLVM_jll
using Ncurses_jll
using Zlib_jll
using libmicrohttpd_jll
using libsndfile_jll
JLLWrappers.@generate_wrapper_header("faust")
JLLWrappers.@declare_executable_product(faust)
JLLWrappers.@declare_file_product(faust2alqt)
JLLWrappers.@declare_file_product(faust2alsa)
JLLWrappers.@declare_file_product(faust2alsaconsole)
JLLWrappers.@declare_file_product(faust2android)
JLLWrappers.@declare_file_product(faust2androidunity)
JLLWrappers.@declare_file_product(faust2api)
JLLWrappers.@declare_file_product(faust2atomsnippets)
JLLWrappers.@declare_file_product(faust2au)
JLLWrappers.@declare_file_product(faust2bela)
JLLWrappers.@declare_file_product(faust2caqt)
JLLWrappers.@declare_file_product(faust2caqtios)
JLLWrappers.@declare_file_product(faust2csound)
JLLWrappers.@declare_file_product(faust2csvplot)
JLLWrappers.@declare_file_product(faust2dplug)
JLLWrappers.@declare_file_product(faust2dssi)
JLLWrappers.@declare_file_product(faust2dummy)
JLLWrappers.@declare_file_product(faust2dummymem)
JLLWrappers.@declare_file_product(faust2eps)
JLLWrappers.@declare_file_product(faust2esp32)
JLLWrappers.@declare_file_product(faust2faustvst)
JLLWrappers.@declare_file_product(faust2firefox)
JLLWrappers.@declare_file_product(faust2gen)
JLLWrappers.@declare_file_product(faust2graph)
JLLWrappers.@declare_file_product(faust2graphviewer)
JLLWrappers.@declare_file_product(faust2ios)
JLLWrappers.@declare_file_product(faust2jack)
JLLWrappers.@declare_file_product(faust2jackconsole)
JLLWrappers.@declare_file_product(faust2jackrust)
JLLWrappers.@declare_file_product(faust2jackserver)
JLLWrappers.@declare_file_product(faust2jaqt)
JLLWrappers.@declare_file_product(faust2jaqtchain)
JLLWrappers.@declare_file_product(faust2juce)
JLLWrappers.@declare_file_product(faust2ladspa)
JLLWrappers.@declare_file_product(faust2linuxunity)
JLLWrappers.@declare_file_product(faust2lv2)
JLLWrappers.@declare_file_product(faust2mathdoc)
JLLWrappers.@declare_file_product(faust2mathviewer)
JLLWrappers.@declare_file_product(faust2max6)
JLLWrappers.@declare_file_product(faust2md)
JLLWrappers.@declare_file_product(faust2msp)
JLLWrappers.@declare_file_product(faust2netjackconsole)
JLLWrappers.@declare_file_product(faust2netjackqt)
JLLWrappers.@declare_file_product(faust2nodejs)
JLLWrappers.@declare_file_product(faust2octave)
JLLWrappers.@declare_file_product(faust2osxiosunity)
JLLWrappers.@declare_file_product(faust2owl)
JLLWrappers.@declare_file_product(faust2paqt)
JLLWrappers.@declare_file_product(faust2pdf)
JLLWrappers.@declare_file_product(faust2plot)
JLLWrappers.@declare_file_product(faust2png)
JLLWrappers.@declare_file_product(faust2portaudiorust)
JLLWrappers.@declare_file_product(faust2pure)
JLLWrappers.@declare_file_product(faust2puredata)
JLLWrappers.@declare_file_product(faust2raqt)
JLLWrappers.@declare_file_product(faust2ros)
JLLWrappers.@declare_file_product(faust2rosgtk)
JLLWrappers.@declare_file_product(faust2rpialsaconsole)
JLLWrappers.@declare_file_product(faust2rpinetjackconsole)
JLLWrappers.@declare_file_product(faust2sam)
JLLWrappers.@declare_file_product(faust2sc)
JLLWrappers.@declare_file_product(faust2sig)
JLLWrappers.@declare_file_product(faust2sigviewer)
JLLWrappers.@declare_file_product(faust2smartkeyb)
JLLWrappers.@declare_file_product(faust2sndfile)
JLLWrappers.@declare_file_product(faust2soul)
JLLWrappers.@declare_file_product(faust2supercollider)
JLLWrappers.@declare_file_product(faust2svg)
JLLWrappers.@declare_file_product(faust2svgviewer)
JLLWrappers.@declare_file_product(faust2teensy)
JLLWrappers.@declare_file_product(faust2unity)
JLLWrappers.@declare_file_product(faust2vcvrack)
JLLWrappers.@declare_file_product(faust2vst)
JLLWrappers.@declare_file_product(faust2vsti)
JLLWrappers.@declare_file_product(faust2w32max6)
JLLWrappers.@declare_file_product(faust2w32msp)
JLLWrappers.@declare_file_product(faust2w32puredata)
JLLWrappers.@declare_file_product(faust2w32vst)
JLLWrappers.@declare_file_product(faust2w64max6)
JLLWrappers.@declare_file_product(faust2w64vst)
JLLWrappers.@declare_file_product(faust2wasm)
JLLWrappers.@declare_file_product(faust2webaudiowasm)
JLLWrappers.@declare_file_product(faust2webaudiowast)
JLLWrappers.@declare_file_product(faust2winunity)
JLLWrappers.@declare_file_product(faustoptflags)
JLLWrappers.@declare_library_product(libHTTPDFaust, "libHTTPDFaust.so.0")
JLLWrappers.@declare_library_product(libOSCFaust, "libOSCFaust.so.1")
JLLWrappers.@declare_library_product(libfaust, "libfaust.so.2")
JLLWrappers.@declare_file_product(stdfaust_lib)
JLLWrappers.@declare_file_product(usage_sh)
function __init__()
    JLLWrappers.@generate_init_header(LLVM_jll, Ncurses_jll, Zlib_jll, libmicrohttpd_jll, libsndfile_jll)
    JLLWrappers.@init_executable_product(
        faust,
        "bin/faust",
    )

    JLLWrappers.@init_file_product(
        faust2alqt,
        "bin/faust2alqt",
    )

    JLLWrappers.@init_file_product(
        faust2alsa,
        "bin/faust2alsa",
    )

    JLLWrappers.@init_file_product(
        faust2alsaconsole,
        "bin/faust2alsaconsole",
    )

    JLLWrappers.@init_file_product(
        faust2android,
        "bin/faust2android",
    )

    JLLWrappers.@init_file_product(
        faust2androidunity,
        "bin/faust2androidunity",
    )

    JLLWrappers.@init_file_product(
        faust2api,
        "bin/faust2api",
    )

    JLLWrappers.@init_file_product(
        faust2atomsnippets,
        "bin/faust2atomsnippets",
    )

    JLLWrappers.@init_file_product(
        faust2au,
        "bin/faust2au",
    )

    JLLWrappers.@init_file_product(
        faust2bela,
        "bin/faust2bela",
    )

    JLLWrappers.@init_file_product(
        faust2caqt,
        "bin/faust2caqt",
    )

    JLLWrappers.@init_file_product(
        faust2caqtios,
        "bin/faust2caqtios",
    )

    JLLWrappers.@init_file_product(
        faust2csound,
        "bin/faust2csound",
    )

    JLLWrappers.@init_file_product(
        faust2csvplot,
        "bin/faust2csvplot",
    )

    JLLWrappers.@init_file_product(
        faust2dplug,
        "bin/faust2dplug",
    )

    JLLWrappers.@init_file_product(
        faust2dssi,
        "bin/faust2dssi",
    )

    JLLWrappers.@init_file_product(
        faust2dummy,
        "bin/faust2dummy",
    )

    JLLWrappers.@init_file_product(
        faust2dummymem,
        "bin/faust2dummymem",
    )

    JLLWrappers.@init_file_product(
        faust2eps,
        "bin/faust2eps",
    )

    JLLWrappers.@init_file_product(
        faust2esp32,
        "bin/faust2esp32",
    )

    JLLWrappers.@init_file_product(
        faust2faustvst,
        "bin/faust2faustvst",
    )

    JLLWrappers.@init_file_product(
        faust2firefox,
        "bin/faust2firefox",
    )

    JLLWrappers.@init_file_product(
        faust2gen,
        "bin/faust2gen",
    )

    JLLWrappers.@init_file_product(
        faust2graph,
        "bin/faust2graph",
    )

    JLLWrappers.@init_file_product(
        faust2graphviewer,
        "bin/faust2graphviewer",
    )

    JLLWrappers.@init_file_product(
        faust2ios,
        "bin/faust2ios",
    )

    JLLWrappers.@init_file_product(
        faust2jack,
        "bin/faust2jack",
    )

    JLLWrappers.@init_file_product(
        faust2jackconsole,
        "bin/faust2jackconsole",
    )

    JLLWrappers.@init_file_product(
        faust2jackrust,
        "bin/faust2jackrust",
    )

    JLLWrappers.@init_file_product(
        faust2jackserver,
        "bin/faust2jackserver",
    )

    JLLWrappers.@init_file_product(
        faust2jaqt,
        "bin/faust2jaqt",
    )

    JLLWrappers.@init_file_product(
        faust2jaqtchain,
        "bin/faust2jaqtchain",
    )

    JLLWrappers.@init_file_product(
        faust2juce,
        "bin/faust2juce",
    )

    JLLWrappers.@init_file_product(
        faust2ladspa,
        "bin/faust2ladspa",
    )

    JLLWrappers.@init_file_product(
        faust2linuxunity,
        "bin/faust2linuxunity",
    )

    JLLWrappers.@init_file_product(
        faust2lv2,
        "bin/faust2lv2",
    )

    JLLWrappers.@init_file_product(
        faust2mathdoc,
        "bin/faust2mathdoc",
    )

    JLLWrappers.@init_file_product(
        faust2mathviewer,
        "bin/faust2mathviewer",
    )

    JLLWrappers.@init_file_product(
        faust2max6,
        "bin/faust2max6",
    )

    JLLWrappers.@init_file_product(
        faust2md,
        "bin/faust2md",
    )

    JLLWrappers.@init_file_product(
        faust2msp,
        "bin/faust2msp",
    )

    JLLWrappers.@init_file_product(
        faust2netjackconsole,
        "bin/faust2netjackconsole",
    )

    JLLWrappers.@init_file_product(
        faust2netjackqt,
        "bin/faust2netjackqt",
    )

    JLLWrappers.@init_file_product(
        faust2nodejs,
        "bin/faust2nodejs",
    )

    JLLWrappers.@init_file_product(
        faust2octave,
        "bin/faust2octave",
    )

    JLLWrappers.@init_file_product(
        faust2osxiosunity,
        "bin/faust2osxiosunity",
    )

    JLLWrappers.@init_file_product(
        faust2owl,
        "bin/faust2owl",
    )

    JLLWrappers.@init_file_product(
        faust2paqt,
        "bin/faust2paqt",
    )

    JLLWrappers.@init_file_product(
        faust2pdf,
        "bin/faust2pdf",
    )

    JLLWrappers.@init_file_product(
        faust2plot,
        "bin/faust2plot",
    )

    JLLWrappers.@init_file_product(
        faust2png,
        "bin/faust2png",
    )

    JLLWrappers.@init_file_product(
        faust2portaudiorust,
        "bin/faust2portaudiorust",
    )

    JLLWrappers.@init_file_product(
        faust2pure,
        "bin/faust2pure",
    )

    JLLWrappers.@init_file_product(
        faust2puredata,
        "bin/faust2puredata",
    )

    JLLWrappers.@init_file_product(
        faust2raqt,
        "bin/faust2raqt",
    )

    JLLWrappers.@init_file_product(
        faust2ros,
        "bin/faust2ros",
    )

    JLLWrappers.@init_file_product(
        faust2rosgtk,
        "bin/faust2rosgtk",
    )

    JLLWrappers.@init_file_product(
        faust2rpialsaconsole,
        "bin/faust2rpialsaconsole",
    )

    JLLWrappers.@init_file_product(
        faust2rpinetjackconsole,
        "bin/faust2rpinetjackconsole",
    )

    JLLWrappers.@init_file_product(
        faust2sam,
        "bin/faust2sam",
    )

    JLLWrappers.@init_file_product(
        faust2sc,
        "bin/faust2sc",
    )

    JLLWrappers.@init_file_product(
        faust2sig,
        "bin/faust2sig",
    )

    JLLWrappers.@init_file_product(
        faust2sigviewer,
        "bin/faust2sigviewer",
    )

    JLLWrappers.@init_file_product(
        faust2smartkeyb,
        "bin/faust2smartkeyb",
    )

    JLLWrappers.@init_file_product(
        faust2sndfile,
        "bin/faust2sndfile",
    )

    JLLWrappers.@init_file_product(
        faust2soul,
        "bin/faust2soul",
    )

    JLLWrappers.@init_file_product(
        faust2supercollider,
        "bin/faust2supercollider",
    )

    JLLWrappers.@init_file_product(
        faust2svg,
        "bin/faust2svg",
    )

    JLLWrappers.@init_file_product(
        faust2svgviewer,
        "bin/faust2svgviewer",
    )

    JLLWrappers.@init_file_product(
        faust2teensy,
        "bin/faust2teensy",
    )

    JLLWrappers.@init_file_product(
        faust2unity,
        "bin/faust2unity",
    )

    JLLWrappers.@init_file_product(
        faust2vcvrack,
        "bin/faust2vcvrack",
    )

    JLLWrappers.@init_file_product(
        faust2vst,
        "bin/faust2vst",
    )

    JLLWrappers.@init_file_product(
        faust2vsti,
        "bin/faust2vsti",
    )

    JLLWrappers.@init_file_product(
        faust2w32max6,
        "bin/faust2w32max6",
    )

    JLLWrappers.@init_file_product(
        faust2w32msp,
        "bin/faust2w32msp",
    )

    JLLWrappers.@init_file_product(
        faust2w32puredata,
        "bin/faust2w32puredata",
    )

    JLLWrappers.@init_file_product(
        faust2w32vst,
        "bin/faust2w32vst",
    )

    JLLWrappers.@init_file_product(
        faust2w64max6,
        "bin/faust2w64max6",
    )

    JLLWrappers.@init_file_product(
        faust2w64vst,
        "bin/faust2w64vst",
    )

    JLLWrappers.@init_file_product(
        faust2wasm,
        "bin/faust2wasm",
    )

    JLLWrappers.@init_file_product(
        faust2webaudiowasm,
        "bin/faust2webaudiowasm",
    )

    JLLWrappers.@init_file_product(
        faust2webaudiowast,
        "bin/faust2webaudiowast",
    )

    JLLWrappers.@init_file_product(
        faust2winunity,
        "bin/faust2winunity",
    )

    JLLWrappers.@init_file_product(
        faustoptflags,
        "bin/faustoptflags",
    )

    JLLWrappers.@init_library_product(
        libHTTPDFaust,
        "lib/libHTTPDFaust.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libOSCFaust,
        "lib/libOSCFaust.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libfaust,
        "lib/libfaust.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_file_product(
        stdfaust_lib,
        "share/faust/stdfaust.lib",
    )

    JLLWrappers.@init_file_product(
        usage_sh,
        "bin/usage.sh",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text: function jacobian_ad(m::HetDSGEGovDebt)


text:
# __precompile__()
module CooperationPolarization

# structs in setup.jl
export Sets, Game, Simulation, Population
export Tracker

# functions in setup.jl
export compute_total_memberships
export party_assignment
export set_members_and_pairs
export random_sets
export vectobase3

# functions in simulate.jl
export evolve!
export update_payoffs_using_matrix!
export update_strategies_and_opinions_db!
export update_opinions!
export update_strategies!
export no_strategy_or_opinion_update!
export compute_payoffs!

# functions in plot.jl
export plot_strat_dist
export plot_opn_dist
export plot_set_dist
export plot_total_opn
export base3tovec
export plotter

# functions in track.jl
export track!
export track_payoff_means!
export track_strat_dist!
export track_pop_dist!
export track_running_strat_dist!
export track_opn_dist!
export track_total_means_and_vars!
export track_total_cooperation!
export track_opn_simpson!
export track_sets_simpson!
export track_cityblock!
export track_hamming!
export make_summary_table
export make_col_list
export make_ens_summary

# functions in simulate_neutral.jl
export evolve_neutral!
export update_actions_neutral!
export compute_actions_neutral!
export update_strategies_and_opinions_neutral!

# struct and functions in track_neutral.jl
export NeutralTracker
export track_neutral!

# load
include("setup.jl")
include("simulate.jl")
include("track.jl")
include("plot.jl")
include("simulate_neutral.jl")
include("track_neutral.jl")

end


text:
__precompile__()

module ShaleDrillingModel

using Distributed
using StatsFuns
using StatsBase
using AxisAlgorithms
using Calculus
using GenGlobal
using Interpolations

using SharedArrays
using SparseArrays
using LinearAlgebra

const AbstractArray3{T} = AbstractArray{T,3}
const AbstractArray4{T} = AbstractArray{T,4}
const AbstractArray5{T} = AbstractArray{T,5}

include("state-space.jl")
include("helpers.jl")
include("learning_transition.jl")

include("flow-payoffs.jl")

include("vf_structs.jl")
include("utility.jl")

include("logsumexp3.jl")

include("makeIminusTVp.jl")

include("vfit.jl")

include("vf_solve_terminal_and_infill.jl")
include("learning_update.jl")
include("vf_solve_exploratory.jl")
include("vf_solve_all.jl")
include("check_dEV.jl")

include("parallel_solution.jl")
include("BSplineExtensions.jl")

include("vf_interpolation.jl")

include("action_probabilities.jl")
include("prob_drill_and_landowner_rev.jl")

# end module
end


text:
module QPointCellFields

using Gridap.Geometry
using Gridap.Fields
using Gridap.ReferenceFEs
using Gridap.CellData

domain = (0,1,0,1)
n = 3
partition = (n,n)
model = CartesianDiscreteModel(domain,partition)

trian = Triangulation(model)
degree = 4
quad = CellQuadrature(trian,degree)
q = get_coordinates(quad)

r = CellField(0.0,trian,quad)
r_q = evaluate(r,q)
test_cell_field(r,q,r_q)

end # module


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

name = "OpenLSTO"
version = v"1.0"

# Collection of sources required to complete build
sources = [
    GitSource("https://github.com/M2DOLab/OpenLSTO.git", "853b9cc433b84c60cd1318a36f9a1d0e6cf65877"),
    DirectorySource("./bundled")
]

# Bash recipe for building across all platforms
# New makefiles added, the patches fix some weird include issues mostly.
# There is likely a better way to fix them, or upstream the fixes.
script = raw"""
cp makefile_FEA $WORKSPACE/srcdir/OpenLSTO/M2DO_FEA/makefile
cp makefile_LSM $WORKSPACE/srcdir/OpenLSTO/M2DO_LSM/makefile
cp makefile_3D_LSM $WORKSPACE/srcdir/OpenLSTO/M2DO_3D_LSM/makefile
# At the moment we need GCC to build with OpenMP
export CXX=g++
cd $WORKSPACE/srcdir/OpenLSTO
atomic_patch -p1 ../patches/include_and_eigenpath.patch
cd $WORKSPACE/srcdir/OpenLSTO/M2DO_FEA
make -j${nproc}
make install
cd ../M2DO_LSM
mkdir bin
make -j${nproc}
make install
cd ../M2DO_3D_LSM
make -j${nproc}
make install
install_license ${WORKSPACE}/srcdir/OpenLSTO/LICENSE
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = expand_cxxstring_abis(supported_platforms(; experimental=true); skip=p->false)

# The products that we will ensure are always built
products = [
    LibraryProduct("m2do_fea", :m2do_fea),
    LibraryProduct("m2do_lsm", :m2do_lsm),
    LibraryProduct("m2do_3d_lsm", :m2do_3d_lsm)
]

# Dependencies that must be installed before this package can be built
dependencies = [
    BuildDependency("Eigen_jll"),
    Dependency("CompilerSupportLibraries_jll")
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies; preferred_gcc_version=v"6", julia_compat="1.6")

max_stars_repo_name: sharanry/Yggdrasil

max_issues_repo_name: sharanry/Yggdrasil

max_forks_repo_name: sharanry/Yggdrasil


text:
using ExamplePackage
using Test

@testset "ExamplePackage.jl" begin
  # Write your own tests here.
  @test foo() ≈ 0.11388071406436832
  @test foo(1, 1.5) ≈ 0.2731856314283442
  @test_broken foo(1, 0) # tells us this is broken
end


text:
# Determine the interactions between atom pairs in a PDB file

export
    distances,
    bonds,
    angles,
    divalents,
    isring,
    isdbonefour,
    isomonefour,
    isphipsi,
    istightphipsi,
    isloosephipsi,
    issecstr,
    issaltbridge,
    ishbond,
    ishydrophobic,
    interactioninfo,
    findinteractions,
    interactions

"""
Calculate the distance matrix for a list of atoms.
Returns a 2D array of distances.
"""
function distances(atoms::Array{Atom,1})
    n_atoms = length(atoms)
    @assert n_atoms > 0 "No atoms in atom list"
    dists = zeros(n_atoms, n_atoms)
    for i in 1:n_atoms
        for j in 1:i-1
            dist = norm(atoms[i].coords - atoms[j].coords)
            dists[i,j] = dist
            dists[j,i] = dist
        end
    end
    return dists
end

"""
Calculate the bonding matrix for a list of atoms.
Returns a 2D boolean array of bonding.
"""
function bonds(atoms::Array{Atom,1}, bonded_pairs::Array{BondedPair,1})
    n_atoms = length(atoms)
    @assert n_atoms > 0 "No atoms in atom list"
    n_cov_dics = length(bonded_pairs)
    @assert n_cov_dics > 0 "No bonds in bonding dictionary"
    bs = falses(n_atoms, n_atoms)
    n_bonds = 0
    for i in 1:n_atoms
        for j in 1:i-1
            bonded = false
            # Check if the atoms are in the same residue
            if atoms[i].res_n == atoms[j].res_n && atoms[i].chain_id == atoms[j].chain_id
                # Check if the atoms appear as a bond in the bonding dictionary
                for k in 1:n_cov_dics
                    if atoms[i].res_name == bonded_pairs[k].residue
                        if atoms[i].atom_name == bonded_pairs[k].atom_one && atoms[j].atom_name == bonded_pairs[k].atom_two
                            bonded = true
                            break
                        elseif atoms[i].atom_name == bonded_pairs[k].atom_two && atoms[j].atom_name == bonded_pairs[k].atom_one
                            bonded = true
                            break
                        end
                    end
                end
            # Check if the atoms are in adjacent residues and part of the amide bond
            elseif atoms[i].chain_id == atoms[j].chain_id && abs(atoms[i].res_n - atoms[j].res_n) == 1
                if atoms[i].res_n > atoms[j].res_n && atoms[i].atom_name == "N" && atoms[j].atom_name == "C"
                    bonded = true
                elseif atoms[j].res_n > atoms[i].res_n && atoms[i].atom_name == "C" && atoms[j].atom_name == "N"
                    bonded = true
                end
            end
            bs[i,j] = bonded
            bs[j,i] = bonded
            if bonded
                n_bonds += 1
            end
        end
    end
    return bs
end

"""
Calculate the bond angle matrix for a list of atoms.
Returns a 2D boolean array of whether bond angles are present.
"""
function angles(atoms::Array{Atom,1}, bs::BitArray)
    n_atoms = length(atoms)
    @assert n_atoms > 0 "No atoms in atom list"
    @assert n_atoms == size(bs, 1) "Number of atoms in atom list and bonding matrix are not the same"
    angs = falses(n_atoms, n_atoms)
    n_angles = 0
    for i in 1:n_atoms
        for j in 1:i-1
            # Atoms have to be on the same or adjacent residues
            if atoms[i].chain_id == atoms[j].chain_id && abs(atoms[i].res_n - atoms[j].res_n) <= 1
                # Covalently-bonded atoms cannot be part of a 1-3 system
                if !bs[i, j]
                    for k in 1:n_atoms
                        if bs[i, k] && bs[k, j]
                            angs[i, j] = true
                            angs[j, i] = true
                            n_angles += 1
                            break
                        end
                    end
                end
            end
        end
    end
    return angs
end

"""
Calculate the divalent matrix for a list of atoms.
Returns a 2D boolean array of whether divalents are present.
"""
function divalents(atoms::Array{Atom,1}, bs::BitArray, angs::BitArray)
    n_atoms = length(atoms)
    @assert n_atoms > 0 "No atoms in atom list"
    @assert n_atoms == size(bs, 1) "Number of atoms in atom list and bonding matrix are not the same"
    @assert n_atoms == size(angs, 1) "Number of atoms in atom list and angle matrix are not the same"
    divs = falses(n_atoms, n_atoms)
    n_divalents = 0
    for i in 1:n_atoms
        for j in 1:i-1
            # Atoms have to be on the same or adjacent residues
            if atoms[i].chain_id == atoms[j].chain_id && abs(atoms[i].res_n - atoms[j].res_n) <= 1
                # Being part of a covalent or 1-3 system prevents a pair from being part of a 1-4 system
                if !bs[i, j] && !angs[i, j]
                    for k in 1:n_atoms
                        if angs[i, k] && bs[k, j]
                            divs[i, j] = true
                            divs[j, i] = true
                            n_divalents += 1
                            break
                        end
                    end
                end
            end
        end
    end
    return divs
end

"""
Determines if two atoms are part of the same ring system.
"""
function isring(atom_one::Atom, atom_two::Atom)
    answer = false
    # See if the atoms are in the same residue
    if atom_one.res_n == atom_two.res_n && atom_one.chain_id == atom_two.chain_id
        if atom_one.res_name == "HIS"
            ring_res = ["CD2", "CG", "CE1", "ND1", "NE2", "1HD1", "1HD2", "1HE1", "1HE2"]
        elseif atom_one.res_name == "PHE"
            ring_res = ["CD2", "CG", "CZ", "CD1", "CE1", "CE2", "HD1", "HD2", "HE1", "HE2", "HZ"]
        elseif atom_one.res_name == "TYR"
            ring_res = ["CD2", "CG", "CZ", "CD1", "CE1", "CE2", "HD1", "HD2", "HE1", "HE2"]
        elseif atom_one.res_name == "TRP"
            ring_res = ["CZ2", "CZ3", "CG", "CH2", "CE3", "CE2", "CD2", "NE1", "CD1", "HD1", "HE1", "HE3", "HZ2", "HZ3", "HH2"]
        else
            ring_res = []
        end
        if atom_one.atom_name in ring_res && atom_two.atom_name in ring_res
            answer = true
        end
    end
    return answer
end

"""
Determines if two atoms are a 1-4 pair around a restricted side chain double bond.
Assumes atoms are a 1-4 pair.
"""
function isdbonefour(atom_one::Atom, atom_two::Atom)
    answer = false
    # See if the atoms are in the same residue
    if atom_one.res_n == atom_two.res_n && atom_one.chain_id == atom_two.chain_id
        if atom_one.res_name == "ASN"
            rest_res = ["CB", "OD1", "1HD1", "2HD1"]
        elseif atom_one.res_name == "GLN"
            rest_res = ["CG", "OE1", "1HE2", "2HE2"]
        elseif atom_one.res_name == "ARG"
            rest_res = ["CD", "NE", "NH1", "NH2", "HE", "1HH1", "2HH1", "1HH2", "2HH2"]
        else
            rest_res = []
        end
        if atom_one.atom_name in rest_res && atom_two.atom_name in rest_res
            answer = true
        end
    end
    return answer
end

"""
Determines if two atoms are a 1-4 pair around the omega angle.
Assumes atoms are a 1-4 pair.
"""
function isomonefour(atom_one::Atom, atom_two::Atom)
    om_res = ["CA", "O", "H"]
    if atom_one.atom_name in om_res && atom_two.atom_name in om_res
        answer = true
    else
        answer = false
    end
    return answer
end

"""
Determines if two atoms are a 1-4 pair around the phi/psi angle.
Assumes atoms are a 1-4 pair.
Glycine H is not polar so should not be present.
"""
function isphipsi(atom_one::Atom, atom_two::Atom)
    pp_res = ["N", "O", "CB", "C", "H"]
    if atom_one.atom_name in pp_res && atom_two.atom_name in pp_res
        answer = true
    else
        answer = false
    end
    return answer
end

"""
Determines if two atoms in a 1-4 pair around the phi/psi angle are either in the same helix/strand or one of the residues is proline.
Assumes atoms are a 1-4 pair around the phi/psi angle.
"""
function istightphipsi(atom_one::Atom, atom_two::Atom, dssps::Dict{String, Char})
    answer = false
    if atom_one.res_name == "PRO" || atom_two.res_name == "PRO"
        answer = true
    else
        dssp_one = dssps["$(atom_one.res_n)$(atom_one.chain_id)"]
        dssp_two = dssps["$(atom_two.res_n)$(atom_two.chain_id)"]
        if dssp_one == dssp_two && (dssp_one == 'H' || dssp_one == 'E')
            answer = true
        end
    end
    return answer
end

"""
Determines if two atoms in a 1-4 pair around the phi/psi angle are either in a loop region or one of the residues is glycine.
Assumes atoms are a 1-4 pair around the phi/psi angle.
"""
function isloosephipsi(atom_one::Atom, atom_two::Atom, dssps::Dict{String, Char})
    answer = false
    if atom_one.res_name == "GLY" || atom_two.res_name == "GLY"
        answer = true
    else
        dssp_one = dssps["$(atom_one.res_n)$(atom_one.chain_id)"]
        dssp_two = dssps["$(atom_two.res_n)$(atom_two.chain_id)"]
        if dssp_one == dssp_two && dssp_one == ' '
            answer = true
        end
    end
    return answer
end

"""
Determines if two atoms are backbone atoms that are part of the same helix/strand and not more than 4 residues apart.
Does not include 3-10 and pi helices.
"""
function issecstr(atom_one::Atom, atom_two::Atom, atoms::Array{Atom,1}, dssps::Dict{String, Char})
    answer = false
    if abs(atom_one.res_n - atom_two.res_n) <= 4 && atom_one.chain_id == atom_two.chain_id
        # See if the atoms have the same secondary structure that is a helix/strand
        dssp_one = dssps["$(atom_one.res_n)$(atom_one.chain_id)"]
        dssp_two = dssps["$(atom_two.res_n)$(atom_two.chain_id)"]
        if dssp_one == dssp_two && (dssp_one == 'H' || dssp_one == 'E')
            # See if the atoms are part of the backbone - currently not counting amide H
            bb_atoms = ["CA", "C", "O", "N"]
            if atom_one.atom_name in bb_atoms && atom_two.atom_name in bb_atoms
                answer = true
            end
            # Check residues between are in the same helix/strand
            for i in 1:length(atoms)
                if atoms[i].res_n > min(atom_one.res_n, atom_two.res_n) && atoms[i].res_n < max(atom_one.res_n, atom_two.res_n) && dssps["$(atoms[i].res_n)$(atoms[i].chain_id)"] != dssp_one && atoms[i].chain_id == atom_one.chain_id
                    answer = false
                    break
                end
            end
        end
    end
    return answer
end

"""
Determines if two atoms are part of oppositely charged groups that are within 4 Angstroms.
His not included.
"""
function issaltbridge(atom_one::Atom, atom_two::Atom, atoms::Array{Atom,1}, dists::Array{Float64})
    answer = false
    # See if the atoms are part of oppositely-charged residues
    if (atom_one.res_name == "ARG" || atom_one.res_name == "LYS") && (atom_two.res_name == "ASP" || atom_two.res_name == "GLU")
        atom_pos = atom_one
        atom_neg = atom_two
    elseif (atom_one.res_name == "ASP" || atom_one.res_name == "GLU") && (atom_two.res_name == "ARG" || atom_two.res_name == "LYS")
        atom_neg = atom_one
        atom_pos = atom_two
    else
        return false
    end
    # See if the atoms are found in the salt bridge parts of their residues
    if atom_pos.res_name == "ARG"
        opts_pos = ["CZ", "NH1", "NH2", "1HH1", "2HH1", "1HH2", "2HH2"]
    elseif atom_pos.res_name == "LYS"
        opts_pos = ["NZ"]
    end
    if atom_neg.res_name == "ASP"
        opts_neg = ["CG", "OD1", "OD2"]
    elseif atom_neg.res_name == "GLU"
        opts_neg = ["CD", "OE1", "OE2"]
    end
    if atom_pos.atom_name in opts_pos && atom_neg.atom_name in opts_neg
        # See if the closest interatomic distance between the positive and negative parts is close enough
        n_atoms = length(atoms)
        ind_pos = Int[]
        ind_neg = Int[]
        for i in 1:n_atoms
            if atom_pos.res_n == atoms[i].res_n && atom_pos.chain_id == atoms[i].chain_id && atoms[i].atom_name in opts_pos
                push!(ind_pos, i)
            elseif atom_neg.res_n == atoms[i].res_n && atom_neg.chain_id == atoms[i].chain_id && atoms[i].atom_name in opts_neg
                push!(ind_neg, i)
            end
        end
        for i in ind_pos
            for j in ind_neg
                if dists[i,j] < 4.0
                    answer = true
                    break
                end
            end
        end
    end
    return answer
end

"""
Determines if two atoms are part of a hydrogen bond where the donor-acceptor distance is <= 3.5 Angstroms, the
hydrogen-acceptor distance is <= 2.5 Angstroms and the donor-hydrogen-acceptor angle is minimally 90 degrees.
Won't pick up the donor-H pairing, even though it is part of the H bond, as this pairing is covalent.
Sulfur not considered.
"""
function ishbond(atom_id_one::Integer, atom_id_two::Integer, atoms::Array{Atom,1}, dists::Array{Float64}, bs::BitArray)
    answer = false
    # Apply crude distance limit and check that atoms are not on the same residue
    if dists[atom_id_one, atom_id_two] <= 3.5 && (atoms[atom_id_one].res_n != atoms[atom_id_two].res_n || atoms[atom_id_one].chain_id != atoms[atom_id_two].chain_id)
        pol_res = ["N", "O"]
        n_atoms = length(atoms)
        atom_h_ids = Int[]
        # Crude way of re-checking distance later
        dist = 100
        # Case where the two atoms are the acceptor and donor
        if atoms[atom_id_one].element in pol_res && atoms[atom_id_two].element in pol_res
            atom_id_a = atom_id_one
            atom_id_b = atom_id_two
            # Find all Hs bonded to either atom that are within 2.5 Angstroms of the other atom
            for i in 1:n_atoms
                if atoms[i].element == "H" && ((bs[i,atom_id_one] && dists[i,atom_id_two] <= 2.5) || (bs[i,atom_id_two] && dists[i,atom_id_one] <= 2.5))
                    push!(atom_h_ids, i)
                end
            end
            # This test was already passed above
            dist = 0
        # Case where atom one is acceptor and atom two is H
        elseif dists[atom_id_one,atom_id_two] <= 2.5 && atoms[atom_id_one].element in pol_res && atoms[atom_id_two].element == "H"
            push!(atom_h_ids, atom_id_two)
            atom_id_a = atom_id_one
            for i in 1:n_atoms
                if bs[i,atom_id_two]
                    atom_id_b = i
                    dist = dists[i,atom_id_one]
                    break
                end
            end
        # Case where atom two is acceptor and atom one is H
        elseif dists[atom_id_one,atom_id_two] <= 2.5 && atoms[atom_id_two].element in pol_res && atoms[atom_id_one].element == "H"
            push!(atom_h_ids, atom_id_one)
            atom_id_a = atom_id_two
            for i in 1:n_atoms
                if bs[i,atom_id_one]
                    atom_id_b = i
                    dist = dists[i,atom_id_two]
                    break
                end
            end
        else
            return false
        end
        if dist <= 3.5
            for atom_id_h in atom_h_ids
                vec_a = atoms[atom_id_h].coords - atoms[atom_id_a].coords
                vec_b = atoms[atom_id_h].coords - atoms[atom_id_b].coords
                if dot(vec_a, vec_b) <= 0
                    answer = true
                    break
                end
            end
        end
    end
    return answer
end

"""
Determines if two atoms are closer than the sum of the VDW radii plus extra_dist Angstroms.
Currently counting C and H atoms only.
Applies a crude 5 A distance cutoff for speed, so extra_dist values above ~1.5 A will have no effect.
"""
function ishydrophobic(atom_one::Atom, atom_two::Atom, distance::Real, extra_dist::Real)
    if distance < 5 && distance < (vdw_radius[atom_one.element] + vdw_radius[atom_two.element] + extra_dist)
        hydr_list = ["C", "H"]
        if atom_one.element in hydr_list && atom_two.element in hydr_list
            answer = true
        else
            answer = false
        end
    else
        answer = false
    end
    return answer
end

"""
Returns a list with interaction types and a list with corresponding distance tolerances.
Zero index represents same atom.
`tolerance_weight` weights all tolerances.
"""
function interactioninfo(tolerance_weight::Real=defaults["tolerance_weight"])
    @assert tolerance_weight >= 0.0 "tolerance_weight cannot be negative"
    n_inter_types = 15
    inter_types = repeat([""], inner=[n_inter_types])
    tolerances = zeros(n_inter_types)
    inter_types[1] = "Covalently bonded"
    tolerances[1] = 0.02
    inter_types[2] = "1-3 system"
    tolerances[2] = 0.05
    inter_types[3] = "Same ring system"
    tolerances[3] = 0.1
    inter_types[4] = "Side-chain restricted 1-4"
    tolerances[4] = 0.1
    inter_types[5] = "Omega 1-4"
    tolerances[5] = 0.1
    inter_types[6] = "Tight phi/psi 1-4"
    tolerances[6] = 0.2
    inter_types[7] = "Loose phi/psi 1-4"
    tolerances[7] = 0.4
    inter_types[8] = "Other phi/psi 1-4"
    tolerances[8] = 0.3
    inter_types[9] = "Other 1-4"
    tolerances[9] = 0.4
    inter_types[10] = "Secondary structure"
    tolerances[10] = 0.5
    inter_types[11] = "Salt bridge"
    tolerances[11] = 0.75
    inter_types[12] = "H bond"
    tolerances[12] = 0.5
    inter_types[13] = "Tight hydrophobic"
    tolerances[13] = 0.5
    inter_types[14] = "Loose hydrophobic"
    tolerances[14] = 1.0
    inter_types[15] = "Other pairs"
    tolerances[15] = 5.0
    tolerances *= tolerance_weight
    return inter_types, tolerances
end

"""
Find the interactions between atoms in a list of atoms.
Returns 2D array of interactions by integer label.
Interactions the same as De Groot et. al., Structure, 1997, but cis/trans limits not implemented.
"""
function findinteractions(atoms::Array{Atom,1},
                    dists::Array{Float64},
                    bs::BitArray,
                    angs::BitArray,
                    divs::BitArray,
                    dssps::Dict{String, Char})
    n_atoms = length(atoms)
    @assert n_atoms > 0 "No atoms in atom list"
    @assert n_atoms == size(dists, 1) "Number of atoms in atom list and distance matrix are not the same"
    @assert n_atoms == size(bs, 1) "Number of atoms in atom list and bonding matrix are not the same"
    @assert n_atoms == size(angs, 1) "Number of atoms in atom list and angle matrix are not the same"
    @assert n_atoms == size(divs, 1) "Number of atoms in atom list and divalent matrix are not the same"
    @assert length(dssps) > 0 "DSSP dictionary is empty"
    inters = zeros(Int, n_atoms, n_atoms)
    inter_types, tolerances = interactioninfo()
    inter_counter = zeros(Int, length(inter_types))
    for i in 1:n_atoms
        for j in 1:i-1
            if bs[i,j]
                n_to_set = 1
            elseif angs[i,j]
                n_to_set = 2
            elseif isring(atoms[i], atoms[j])
                n_to_set = 3
            elseif divs[i,j]
                if isdbonefour(atoms[i], atoms[j])
                    n_to_set = 4
                elseif isomonefour(atoms[i], atoms[j])
                    n_to_set = 5
                elseif isphipsi(atoms[i], atoms[j])
                    if istightphipsi(atoms[i], atoms[j], dssps)
                        n_to_set = 6
                    elseif isloosephipsi(atoms[i], atoms[j], dssps)
                        n_to_set = 7
                    else
                        n_to_set = 8
                    end
                else
                    n_to_set = 9
                end
            elseif issecstr(atoms[i], atoms[j], atoms, dssps)
                n_to_set = 10
            elseif issaltbridge(atoms[i], atoms[j], atoms, dists)
                n_to_set = 11
            elseif ishbond(i, j, atoms, dists, bs)
                n_to_set = 12
            elseif ishydrophobic(atoms[i], atoms[j], dists[i, j], 0.5)
                n_to_set = 13
            elseif ishydrophobic(atoms[i], atoms[j], dists[i, j], 1.0)
                n_to_set = 14
            else
                n_to_set = 15
            end
            inters[i,j] = n_to_set
            inters[j,i] = n_to_set
            inter_counter[n_to_set] += 1
        end
    end
    println("Found interactions:")
    max_length = maximum(length.(inter_types))
    for (k, inter_type) in enumerate(inter_types)
        println("  ", rpad(inter_type, max_length), " - ", inter_counter[k])
    end
    if inter_counter[12] == 0
        println("-- NB no H bonds were found - did you remember to add polar hydrogens? --")
    end
    return inters
end

"""
Finds distance constraints from a set of interactions and distances.
Returns a `Constraints` object.
All explicit interactions calculated - other interactions written with probability such that `other_ratio` times n_atoms other interactions are written.
Sparse interactions scale as N but other interactions scale as N^2, hence the use of a ratio
rather than a direct probability to keep the sparse/other ratio roughly constant.
The hard sphere radii are multiplied by 0.8 to reduce the number of erroneous constraints.
"""
function Constraints(atoms::Array{Atom,1},
                    dists::Array{Float64},
                    inters::Array{Int};
                    other_ratio::Real=defaults["other_ratio"],
                    tolerance_weight::Real=defaults["tolerance_weight"])
    n_atoms = length(atoms)
    @assert n_atoms > 0 "No atoms in atom list"
    @assert n_atoms == size(dists, 1) "Number of atoms in atom list and distance matrix are not the same"
    @assert n_atoms == size(inters, 1) "Number of atoms in atom list and interaction matrix are not the same"
    @assert 0.0 <= other_ratio "other_ratio must be positive or zero"
    constraints_lower = Float64[]
    constraints_upper = Float64[]
    inter_types, tolerances = interactioninfo(tolerance_weight)
    present_i = Int[]
    present_j = Int[]
    # Calculate the probability of accepting other interactions based on the desired ratio of other interactions to n_atoms
    other_prob = min((2 * other_ratio) / (n_atoms - 1), 1.0)
    for i in 1:n_atoms
        for j in 1:i-1
            # Calculate explicit interactions and other interactions with a probability other_prob
            if inters[i, j] != 15 || rand() < other_prob
                constraint_upper = dists[i,j] + tolerances[inters[i,j]]
                # Factor of 0.8 used here
                radius_weighting = 0.8
                sum_radii = (hard_sphere_radius[atoms[i].element] + hard_sphere_radius[atoms[j].element]) * radius_weighting
                constraint_lower = max(dists[i,j] - tolerances[inters[i,j]], sum_radii)
                if constraint_lower > constraint_upper
                    println(
                        "Erroneous constraint:\t",
                        atomid(atoms[i]),
                        "\t",
                        atomid(atoms[j]),
                        "\t",
                        round(constraint_lower, digits=3),
                        "\t",
                        round(constraint_upper, digits=3)
                    )
                    constraint_lower = constraint_upper
                end
                push!(constraints_lower, constraint_lower)
                push!(constraints_upper, constraint_upper)
                push!(present_i, i)
                push!(present_j, j)
            end
        end
    end
    println("Found distance constraints from single structure")
    return Constraints(atoms, constraints_lower, constraints_upper, hcat(present_i, present_j))
end

"""
Finds distance constraints from two sets of interactions and distances.
Returns a `Constraints` object.
Input atoms must correspond directly to each other.
"""
function Constraints(atoms::Array{Atom,1},
                    dists_one::Array{Float64},
                    dists_two::Array{Float64},
                    inters_one::Array{Int},
                    inters_two::Array{Int};
                    other_ratio::Real=defaults["other_ratio"],
                    tolerance_weight::Real=defaults["tolerance_weight"])
    n_atoms = length(atoms)
    @assert n_atoms > 0 "No atoms in atom list"
    @assert n_atoms == size(dists_one, 1) "Number of atoms in atom list and distance matrix one are not the same"
    @assert n_atoms == size(dists_two, 1) "Number of atoms in atom list and distance matrix two are not the same"
    @assert n_atoms == size(inters_one, 1) "Number of atoms in atom list and interaction matrix one are not the same"
    @assert n_atoms == size(inters_two, 1) "Number of atoms in atom list and interaction matrix two are not the same"
    @assert 0.0 <= other_ratio "other_ratio must be positive or zero"
    constraints_lower = Float64[]
    constraints_upper = Float64[]
    inter_types, tolerances = interactioninfo(tolerance_weight)
    present_i = Int[]
    present_j = Int[]
    # Calculate the probability of accepting other interactions based on the desired ratio of other interactions to n_atoms
    other_prob = min((2 * other_ratio) / (n_atoms - 1), 1.0)
    for i in 1:n_atoms
        for j in 1:i-1
            # Calculate explicit interactions and other interactions with a probability other_prob
            if (inters_one[i, j] != 15 && inters_two[i, j] != 15) || rand() < other_prob
                constraint_upper_one = dists_one[i,j] + tolerances[inters_one[i,j]]
                constraint_upper_two = dists_two[i,j] + tolerances[inters_two[i,j]]
                constraint_upper = max(constraint_upper_one, constraint_upper_two)
                # Factor of 0.8 used here
                radius_weighting = 0.8
                sum_radii = (hard_sphere_radius[atoms[i].element] + hard_sphere_radius[atoms[j].element]) * radius_weighting
                constraint_lower_one = dists_one[i,j] - tolerances[inters_one[i,j]]
                constraint_lower_two = dists_two[i,j] - tolerances[inters_two[i,j]]
                constraint_lower = max(min(constraint_lower_one, constraint_lower_two), sum_radii)
                if constraint_lower > constraint_upper
                    println(
                        "Erroneous constraint:\t",
                        atomid(atoms[i]),
                        "\t",
                        atomid(atoms[j]),
                        "\t",
                        round(constraint_lower, digits=3),
                        "\t",
                        round(constraint_upper, digits=3)
                    )
                    constraint_lower = constraint_upper
                end
                push!(constraints_lower, constraint_lower)
                push!(constraints_upper, constraint_upper)
                push!(present_i, i)
                push!(present_j, j)
            end
        end
    end
    println("Found combined distance constraints from two structures")
    return Constraints(atoms, constraints_lower, constraints_upper, hcat(present_i, present_j))
end

"""
Run the interaction finding pipeline with one or two structures.
If one structure is given, returns a `Constraints` object.
If two structures are given, returns three `Constraints` objects corresponding to the
combined constraints, constraints from structure one and constraints from structure two.
"""
function interactions(pdb_filepath::AbstractString,
                    dssp_filepath::AbstractString;
                    other_ratio::Real=defaults["other_ratio"],
                    tolerance_weight::Real=defaults["tolerance_weight"])
    atoms = readpdb(pdb_filepath)
    dssps = readdssp(dssp_filepath, atoms)
    dists = distances(atoms)
    bs = bonds(atoms, protein_bonds)
    angs = angles(atoms, bs)
    divs = divalents(atoms, bs, angs)
    inters = findinteractions(atoms, dists, bs, angs, divs, dssps)
    return Constraints(atoms, dists, inters, other_ratio=other_ratio, tolerance_weight=tolerance_weight)
end

function interactions(pdb_filepath_one::AbstractString,
                    dssp_filepath_one::AbstractString,
                    pdb_filepath_two::AbstractString,
                    dssp_filepath_two::AbstractString;
                    other_ratio::Real=defaults["other_ratio"],
                    tolerance_weight::Real=defaults["tolerance_weight"])
    atoms_all_one = readpdb(pdb_filepath_one)
    atoms_all_two = readpdb(pdb_filepath_two)
    atoms_one, atoms_two = findcommonatoms(atoms_all_one, atoms_all_two)
    dssps_one = readdssp(dssp_filepath_one, atoms_one)
    dssps_two = readdssp(dssp_filepath_two, atoms_two)
    dists_one = distances(atoms_one)
    bonds_one = bonds(atoms_one, protein_bonds)
    angles_one = angles(atoms_one, bonds_one)
    divalents_one = divalents(atoms_one, bonds_one, angles_one)
    inters_one = findinteractions(atoms_one, dists_one, bonds_one, angles_one, divalents_one, dssps_one)
    constraints_one = Constraints(atoms_one, dists_one, inters_one, other_ratio=other_ratio, tolerance_weight=tolerance_weight)
    dists_two = distances(atoms_two)
    bonds_two = bonds(atoms_two, protein_bonds)
    angles_two = angles(atoms_two, bonds_two)
    divalents_two = divalents(atoms_two, bonds_two, angles_two)
    inters_two = findinteractions(atoms_two, dists_two, bonds_two, angles_two, divalents_two, dssps_two)
    constraints_two = Constraints(atoms_two, dists_two, inters_two, other_ratio=other_ratio, tolerance_weight=tolerance_weight)
    constraints_com = Constraints(atoms_one, dists_one, dists_two, inters_one, inters_two, other_ratio=other_ratio, tolerance_weight=tolerance_weight)
    return constraints_com, constraints_one, constraints_two
end


text:
struct TreerXgb
    fm
    featur::Vector{Int64}
end

struct TreedaXgb
    fm
    featur::Vector{Int64}
    lev::Vector{Int64}
    ni::Vector{Int64}
end

"""
    treer_xgb(X, y;
        subsample = 1, colsample_bytree = 1, colsample_bynode = 1,
        max_depth = 6, min_child_weight = 5,
        lambda = 0, verbose = false, kwargs...)
Regression tree with XGBoost.
* `X` : X-data (n obs., p variables).
* `y` : Univariate Y-data (n obs.).
* `subsample` : Proportion of rows sampled in `X`
    for building the tree.
* `colsample_bytree` : Proportion of columns sampled in `X`
    for building the tree.
* `colsample_bynode` : Proportion of columns sampled at each node
    in the columns selected for the tree.
* `max_depth` : Maximum depth of the tree.
* `min_child_weight` : Minimum nb. observations that each leaf
    needs to have.
* `lambda` : L2 regularization term on weights.
    Increasing this value will make model more conservative.
* `verbose` : If true, fitting information are printed.
* `kwargs` : Optional named arguments to pass in function `xgboost`
    of `XGBoost.jl` (https://xgboost.readthedocs.io/en/latest/parameter.html).

The function builds a single tree using package `XGboost.jl'.

The sampling of the observations and variables are without replacement.

## References

XGBoost.jl
https://github.com/dmlc/XGBoost.jl

Breiman, L., Friedman, J. H., Olshen, R. A., and Stone, C. J. Classification
And Regression Trees. Chapman & Hall, 1984.

Gey, S., 2002. Bornes de risque, détection de ruptures, boosting :
trois thèmes statistiques autour de CART en régression (These de doctorat).
Paris 11. http://www.theses.fr/2002PA112245
"""
function treer_xgb(X, y;
        subsample = 1, colsample_bytree = 1, colsample_bynode = 1,
        max_depth = 6, min_child_weight = 5,
        lambda = 0, verbose = false, kwargs...)
    X = ensure_mat(X)
    p = size(X, 2)
    num_round = 1
    fm = xgboost(X, num_round; label = Float64.(vec(y)),
        seed = Int64(round(rand(1)[1] * 1e5)),
        booster = :gbtree,
        tree_method = :auto,
        eta = 1, # learning rate
        subsample = subsample,
        colsample_bytree = colsample_bytree, colsample_bylevel = 1,
        colsample_bynode = colsample_bynode,
        max_depth = max_depth, min_child_weight = min_child_weight,
        lambda = lambda,
        silent = !verbose, kwargs...)
    featur = collect(1:p)
    TreerXgb(fm, featur)
end

"""
    rfr_xgb(X, y; rep = 50,
        subsample = .7,
        colsample_bytree = 1, colsample_bynode = 1/3,
        max_depth = 6, min_child_weight = 5,
        lambda = 0, verbose = false, kwargs...)
Random forest regression with XGBoost.
* `X` : X-data (n obs., p variables).
* `y` : Univariate Y-data (n obs.).
* `rep` : Nb. trees to build in the forest.
* `subsample` : Proportion of rows sampled in `X`
    for building each tree.
* `colsample_bytree` : Proportion of columns sampled in `X` at each tree.
* `colsample_bynode` : Proportion of columns sampled at each node
    in the columns selected for the tree.
* `max_depth` : Maximum depth of the trees.
* `min_child_weight` : Minimum nb. observations that each leaf
    needs to have.
* `lambda` : L2 regularization term on weights.
    Increasing this value will make model more conservative.
* `verbose` : If true, fitting information are printed.
* `kwargs` : Optional named arguments to pass in function `xgboost`
    of `XGBoost.jl` (https://xgboost.readthedocs.io/en/latest/parameter.html).

The function uses package `XGboost.jl' to build the forest.
See https://xgboost.readthedocs.io/en/latest/tutorials/rf.html.

## References

XGBoost.jl
https://github.com/dmlc/XGBoost.jl

Breiman, L., 1996. Bagging predictors. Mach Learn 24, 123–140.
https://doi.org/10.1007/BF00058655

Breiman, L., 2001. Random Forests. Machine Learning 45, 5–32.
https://doi.org/10.1023/A:1010933404324

Genuer, R., 2010. Forêts aléatoires : aspects théoriques,
sélection de variables et applications. PhD Thesis. Université Paris Sud - Paris XI.

Gey, S., 2002. Bornes de risque, détection de ruptures, boosting :
trois thèmes statistiques autour de CART en régression (These de doctorat).
Paris 11. http://www.theses.fr/2002PA112245
"""
function rfr_xgb(X, y; rep = 50,
        subsample = .7,
        colsample_bytree = 1, colsample_bynode = 1/3,
        max_depth = 6, min_child_weight = 5,
        lambda = 0, verbose = false, kwargs...)
    X = ensure_mat(X)
    p = size(X, 2)
    num_round = 1
    fm = xgboost(X, num_round; label = Float64.(vec(y)),
        seed = Int64(round(rand(1)[1] * 1e5)),
        booster = :gbtree,
        tree_method = :auto,
        num_parallel_tree = rep,
        eta = 1, # learning rate
        subsample = subsample,
        colsample_bytree = colsample_bytree, colsample_bylevel = 1,
        colsample_bynode = colsample_bynode,
        max_depth = max_depth, min_child_weight = min_child_weight,
        lambda = lambda,
        silent = !verbose, kwargs...)
    featur = collect(1:p)
    TreerXgb(fm, featur)
end

"""
    xgboostr(X, y; rep = 50, eta = .3,
        subsample = .7, colsample_bytree = 1, colsample_bynode = 1/3,
        max_depth = 6, min_child_weight = 5,
        lambda = 1, verbose = false, kwargs...)
XGBoost regression.
* `X` : X-data (n obs., p variables).
* `y` : Univariate Y-data (n obs.).
* `rep` : Nb. trees to build.
* `eta` : Learning rate ([0, 1]).
* `subsample` : Proportion of rows sampled in `X`
    for building each tree.
* `colsample_bytree` : Proportion of columns sampled in `X` at each tree.
* `colsample_bynode` : Proportion of columns sampled at each node
    in the columns selected for the tree.
* `max_depth` : Maximum depth of the trees.
* `min_child_weight` : Minimum nb. observations that each leaf
    needs to have.
* `lambda` : L2 regularization term on weights.
    Increasing this value will make model more conservative.
* `verbose` : If true, fitting information are printed.
* `kwargs` : Optional named arguments to pass in function `xgboost`
    of `XGBoost.jl` (https://xgboost.readthedocs.io/en/latest/parameter.html).

The function uses package `XGboost.jl' to build the forest.
See https://xgboost.readthedocs.io/en/latest/tutorials/rf.html.

## References

XGBoost
https://xgboost.readthedocs.io/en/latest/index.html

XGBoost.jl
https://github.com/dmlc/XGBoost.jl
"""
function xgboostr(X, y; rep = 50, eta = .3,
        subsample = .7, colsample_bytree = 1, colsample_bynode = 1/3,
        max_depth = 6, min_child_weight = 5,
        lambda = 1, verbose = false, kwargs...)
    X = ensure_mat(X)
    p = size(X, 2)
    num_round = rep
    fm = xgboost(X, num_round; label = Float64.(vec(y)),
        seed = Int64(round(rand(1)[1] * 1e5)),
        booster = :gbtree,
        tree_method = :auto,
        num_parallel_tree = 1,
        eta = eta,
        subsample = subsample,
        colsample_bytree = colsample_bytree,
        colsample_bylevel = 1, colsample_bynode = colsample_bynode,
        max_depth = max_depth, min_child_weight = min_child_weight,
        lambda = lambda,
        silent = !verbose, kwargs...)
    featur = collect(1:p)
    TreerXgb(fm, featur)
end

"""
    predict(object::TreerXgb, X)
Compute Y-predictions from a fitted model.
* `object` : The fitted model.
* `X` : X-data for which predictions are computed.
"""
function predict(object::TreerXgb, X)
    X = ensure_mat(X)
    m = size(X, 1)
    pred = XGBoost.predict(object.fm, X)
    pred = reshape(pred, m, 1) ;
    (pred = pred,)
end

"""
    vimp_xgb(object::TreerXgb)
Compute variable (feature) importances from an XGBoost model.
* `object` : The fitted model.

Features with imp = 0 are not returned.
"""
function vimp_xgb(object::Union{TreerXgb, TreedaXgb})
    p = length(object.featur)
    res = XGBoost.importance(object.fm, string.(object.featur))
    fname = [] ; gain = [] ; cover = [] ; freq = []
    for i in res
        push!(fname, i.fname)
        push!(gain, i.gain)
        push!(cover, i.cover)
        push!(freq, i.freq)
    end
    fname = eval(Meta.parse.(fname))
    zgain = zeros(p)
    zcover = copy(zgain)
    zfreq = copy(zgain)
    for i = 1:length(fname)
        s = fname[i]
        zgain[s] = gain[i]
        zcover[s] = cover[i]
        zfreq[s] = freq[i]
    end
    res = DataFrame(hcat(zgain, zcover, zfreq), [:gain, :cover, :freq])
    res.featur = object.featur
    res
end


text:
function _conv2(x, w; pad=0, stride=1, xcorr=false)
    max_pad = map(x->x-1-pad,size(w))
    y = conv2(x, xcorr ? rot180(w) : w)
    return y[1+max_pad[1]:stride:end-max_pad[1], 1+max_pad[2]:stride:end-max_pad[2]]
end

function cudnnConvolutionForward(x::Array{Float32,4}, w, y; padding=0, stride=1,
upscale=1, mode=0, cd=nothing, algorithm="okuru13", workSpace=0, workSpaceSizeInBytes=0, alpha=1, beta=1)
    # x: (W,H,C,N)
    # w: (W,H,C,K)
    # y: (W,H,K,N)
    println("conv on cpu")
    @assert padding==0 && stride==1 && upscale==1&& mode==0&& algorithm=="okuru13"&& workSpace==0&& workSpaceSizeInBytes==0&& alpha==1&& beta==1
    Wx,Hx,Cx,N = size(x)
    Ww,Hw,Cw,K = size(w)
    for n in 1:N, k in 1:K, c in 1:Cx
        y[:,:,k,n] += _conv2(x[:,:,c,n], w[:,:,c,k]; pad=padding, stride=stride, xcorr=mode!=0)
    end
    return y
end

# dw = rot180(xcorr(x,dy))
function cudnnConvolutionBackwardFilter(x::Array{Float32,4}, dy, dw; padding=0, stride=1, upscale=1, mode=0)
    # x:    (Wx,Hx,Cx,N)
    # dy:   (Wy,Hy,K,N)
    # dw:    (Ww,Hw,Cw,K)
    @assert padding==0&& stride==1&& upscale==1&& mode==0
    Wx,Hx,C,Nx = size(x)
    Wy,Hy,K,Ny = size(dy)
    for c in 1:C, k in 1:K, n in 1:Ny
        dw[:,:,c,k] += rot180(_conv2(x[:,:,c,n], dy[:,:,k,n]; pad=padding, stride=stride, xcorr=true))
    end
    return dw
end

# dx = xcorr(dy, w, 'full')
function cudnnConvolutionBackwardData(w::Array{Float32,4}, dy, dx; padding=0, stride=1, upscale=1, mode=0)
    @assert padding==0&& stride==1&& upscale==1&& mode==0
    Wy,Hy,Ky,N = size(dy)
    Ww,Hw,C,Kw = size(w)
    @assert Ky==Kw
    for n in 1:N, c in 1:C, k in 1:Kw
        t = conv2(dy[:,:,k,n], rot180(w[:,:,c,k]))
        # t = _conv2(dy[:,:,k,n], w[:,:,c,k]; pad=Ww-1, stride=stride, xcorr=true)
        dx[:,:,c,n] += t
    end
    return dx
end

function cudnnGetConvolutionNdForwardOutputDim(x::Array{Float32,4},w; padding=padding,stride=stride)
    Wx,Hx,Cx,N = size(x)
    Ww,Hw,Cw,K = size(w)
    @assert Cx==Cw
    Wy,Hy = floor(Int, 1 + (Int[Wx,Hx] + 2*padding - Int[Ww,Hw]) / stride)
end

function cudnnPoolingForward(x::Array{Float32,4}, y; window=2, padding=0, stride=1, mode=0)
    @assert padding==0 && stride==1 && mode==0
    # x: (W,H,C,N)
    Wx,Hx,C,Nx = size(x);
    Wy,Hy,K,Ny = size(y);
    @assert Nx == Ny && C==K
    for n in 1:Nx, c in 1:C, j in 1:Hx-psize+1, i in 1:Wx-psize+1
        y[i,j,c,n] = maximum(x[i:i+psize-1,j:j+psize-1,c,n])
    end
    return y
end

function cudnnPoolingBackward(y::Array{Float32,4}, dy, x, dx; window=2, padding=0, stride=1, mode=0)
    @assert padding==0 && stride==1 && mode==0
    # x: (W,H,C,N)
    Wx,Hx,C,Nx = size(x);
    Wy,Hy,K,Ny = size(y);
    @assert Nx == Ny && C==K
    for n in 1:Nx, c in 1:C, j in 1:Hx-window+1, i in 1:Wx-window+1
        # y2[i,j,c,n] = maximum()
        a = x[i:i+window-1,j:j+window-1,c,n]
        mi,mj = ind2sub(a,indmax(a))
        dx[i+mi-1,j+mj-1,c,n] = dy[i,j,c,n]
    end
    return dx
end

function GetPoolingNdForwardOutputDim(x::Array{Float32,4}; window=2, padding=0, stride=1, mode=0)
    @assert padding==0 && stride==1 && mode==0
    dims = [size(x)...]
    # (mode, pdims, window, padding, stride) = cudnnGetPoolingNdDescriptor(pd)
    for i=1:length(dims)-2
        dims[i] = 1 + ceil((dims[i] + 2*padding - window) / stride)
    end
    tuple(dims...)
end


text:
import Canvas: @api

@api test1 => TestType1 begin
    arg(a::FloatingPoint)
end
@api test1 => TestType1 begin
    arg(a::FloatingPoint)
end
@fact test1(3) => TestType1(3.0)

@api testsub => TestSubType <: Number begin
    arg(a::FloatingPoint)
end
@fact testsub(3) => TestSubType(3.0)

@api test2 => TestType2 begin
    typedarg(a::FloatingPoint)
    arg(b::Int)
    curry(x::Int)
end

@fact_throws test2(2, 3, 4) # Because 2 should be float
@fact test2(2.0, 3, 4) => TestType2(2.0, 3, 4)
@fact test2(2.0, 3)(4) => TestType2(2.0, 3, 4)

@api test3 => TestType3 begin
    typedarg(a::FloatingPoint)
    arg(b::Int)
    curry(x::Int)
    kwarg(p::Any="black")
    typedkwarg(q::FloatingPoint=1)
end

@fact_throws test3(2, 3, 4) # Because 2 should be float
@fact test3(2.0, 3, 4) => TestType3(2.0, 3, 4, "black", 1)
@fact test3(2.0, 3)(4) => TestType3(2.0, 3, 4, "black", 1)
@fact test3(2.0, 3, 4, p=:white) => TestType3(2.0, 3, 4, :white, 1)
@fact test3(2.0, 3, q=5.0)(4) => TestType3(2.0, 3, 4, "black", 5.0)
@fact_throws test3(2.0, 5, 4, q=1)

abstract A
immutable B<:A end

#@show macroexpand(:(
@api test4 => Test4{P <: A} <: Tile begin
    typedarg(a::AbstractArray=Integer[])
    arg(b::P)
    curry(c::Int)
end
#))
println(test4(B())(2))

@api code => Code <: A begin
    arg(language::String="julia")
    arg(code::Any)
end


text:
include("../functions/Functions.jl")
include("../evaluate/Evaluate.jl")
include("../opt/SAGA.jl")
include("../opt/SAGA_PROX.jl")
include("../opt/sgd.jl")
include("../data/data.jl")
include("../plot/Plotting.jl")

using ProximalOperators
using Plotting

X, y = getFullData()

# init initial point x_0
x0 = zeros(size(X,2),1)

lambda = 0.1

iter = 1000
step =0.001

#f = f_logreg(X,y,x0,lambda)
#g = sub_g_logreg()

f = f_svm(X,y,x0,lambda)
g = sub_g_svm()

prox_operator = SqrNormL2(lambda)

@time (theta, strings, iter,val_array_saga) = pgd(X,y,x0,f,g,1e-8,lambda,iter,step,size(X,1),prox_operator)

println(evaluate(X,y,theta,"svm",false))

#@time (theta,state,val_array,stop_array) = sgd(X,y,x0,f,g,1e-8,lambda,iter,step,size(X,1))

plotDoubleArray(val_array_saga,val_array, "saga1.svg")

println(evaluate(X,y,theta,"svm",false))
# Lippschitz = largest eigenvalue of Hessian*Hessian'

max_stars_repo_name: zagazao/num-opt

max_issues_repo_name: zagazao/num-opt

max_forks_repo_name: zagazao/num-opt


text:
# # Barotropic gyre

using Oceananigans.Grids
using Oceananigans.Utils: prettytime, hours, day, days, years

using Statistics
using JLD2
using Printf
using GLMakie

function geographic2cartesian(λ, φ, r=1)
    Nλ = length(λ)
    Nφ = length(φ)

    λ = repeat(reshape(λ, Nλ, 1), 1, Nφ)
    φ = repeat(reshape(φ, 1, Nφ), Nλ, 1)

    λ_azimuthal = λ .+ 180  # Convert to λ ∈ [0°, 360°]
    φ_azimuthal = 90 .- φ   # Convert to φ ∈ [0°, 180°] (0° at north pole)

    x = @. r * cosd(λ_azimuthal) * sind(φ_azimuthal)
    y = @. r * sind(λ_azimuthal) * sind(φ_azimuthal)
    z = @. r * cosd(φ_azimuthal)

    return x, y, z
end

function visualize_barotropic_gyre(filepath)

    file = jldopen(filepath)

    Nx = file["grid/Nx"]
    Ny = file["grid/Ny"]

    # A spherical domain
    grid = RegularLatitudeLongitudeGrid(size = (Nx, Ny, 1),
                                        longitude = (-30, 30),
                                        latitude = (15, 75),
                                        z = (-4000, 0))

    iterations = parse.(Int, keys(file["timeseries/t"]))

    xu, yu, zu = geographic2cartesian(xnodes(Face,   grid), ynodes(Center, grid))
    xv, yv, zv = geographic2cartesian(xnodes(Center, grid), ynodes(Face,   grid))
    xc, yc, zc = geographic2cartesian(xnodes(Center, grid), ynodes(Center, grid))

    iter = Node(0)

    plot_title = @lift @sprintf("Barotropic gyre: time = %s", prettytime(file["timeseries/t/" * string($iter)]))

    u = @lift file["timeseries/u/" * string($iter)][:, :, 1]
    v = @lift file["timeseries/v/" * string($iter)][:, :, 1]
    η = @lift file["timeseries/η/" * string($iter)][:, :, 1]

    fig = Figure(resolution = (2160, 1080))

    x = (xu, xv, xc)
    y = (yu, yv, yc)
    z = (zu, zv, zc)

    statenames = ["u", "v", "η"]
    for (n, var) in enumerate([u, v, η])
        ax = fig[3:7, 3n-2:3n] = LScene(fig) # make plot area wider
        wireframe!(ax, Sphere(Point3f0(0), 0.99f0), show_axis=false)
        surface!(ax, x[n], y[n], z[n], color=var, colormap=:balance) #, colorrange=clims[n])
        rotate_cam!(ax.scene, (3π/4, -π/8, 0))
        zoom!(ax.scene, (0, 0, 0), 5, false)
        fig[2, 2 + 3*(n-1)] = Label(fig, statenames[n], textsize = 50) # put names in center
    end

    supertitle = fig[0, :] = Label(fig, plot_title, textsize=50)

    record(fig, output_prefix * ".mp4", iterations, framerate=30) do i
        @info "Animating iteration $i/$(iterations[end])..."
        iter[] = i
    end

    close(file)

    return nothing
end

#Nx = 720
#Ny = 720

Nx = 360
Ny = 360

output_prefix = "barotropic_gyre_Nx$(Nx)_Ny$(Ny)"
filepath = output_prefix * ".jld2"

visualize_barotropic_gyre(filepath)


text:
ROS.@include std_msgs: Float64MultiArray

sub_callback_enabled = false

function testSubscriber()
    ROS.init("testSub")
    nh = ROS.NodeHandle()

    sub = ROS.subscribe(nh, "test_sub", 1, ROS.std_msgs_Float64MultiArray, callback)

    ROS.getNumPublishers(sub)
    @test ROS.getTopic(sub) == "/test_sub"

    loop = 100
    while loop > 0 && !sub_callback_enabled
        loop -= 1
        ROS.sleep(ROS.Rate(10))
        ROS.spinOnce()
    end
    @test sub_callback_enabled
    ROS.shutdown(sub)
    ROS.shutdown(nh)
    println("All $(basename(@__FILE__)) tests passed.")
end

function callback(t)
    t.data = [1,2,3,4]
    global sub_callback_enabled = true
end


text:
module hw1_4

using LinearAlgebra, PlotlyJS, ORCA

export total_manuever, plot_maneuver

function oe_to_rθh(oe::Vector,μ::Real) :: Vector

  a,e,i,Ω,ω,ν = oe
  return [a*(1-e^2)/(1+e*cos(ν)),
  0,
  0,
  (μ/sqrt(μ*a*(1-e^2)))*e*sin(ν),
  (μ/sqrt(μ*a*(1-e^2)))*(1+e*cos(ν)),
  0]

end

function rθh_to_xyz(rθh_vec::Vector,oe::Vector)

  a,e,i,Ω,ω,ν = oe
  θ = ω+ν
  cΩ,sΩ,ci,si,cθ,sθ = cos(Ω),sin(Ω),cos(i),sin(i),cos(θ),sin(θ)
  DCM = [cΩ*cθ-sΩ*ci*sθ -cΩ*sθ-sΩ*ci*cθ sΩ*si;
  sΩ*cθ+cΩ*ci*sθ -sΩ*sθ+cΩ*ci*cθ -cΩ*si;
  si*sθ si*cθ ci]
  DCM = kron(Matrix(I,2,2),DCM)
  return DCM*rθh_vec

end

function oe_to_xyz(oe::Vector,μ::Real)

  return rθh_to_xyz(oe_to_rθh(oe,μ),oe)

end

function xyz_to_oe(cart_vec::Vector,μ::Real)

  r_xyz, v_xyz = cart_vec[1:3],cart_vec[4:6]
  r = norm(r_xyz)
  h_xyz = cross(r_xyz,v_xyz) #km^2/s
  h = norm(h_xyz) #km^2/s
  ξ = dot(v_xyz,v_xyz)/2 - μ/r #km^2 s^-2
  a = -μ/(2ξ) #km
  e = sqrt(1 + (2h^2*ξ)/μ^2)
  e_xyz = cross(v_xyz,h_xyz)/μ - r_xyz/r
  i = acos(h_xyz[3]/h) #rad
  n_xyz = cross([0,0,1],h_xyz)
  Ω = acos(dot(n_xyz,[1,0,0])/norm(n_xyz))
  ω = acos((dot(n_xyz,e_xyz)/(norm(n_xyz)*e)))
  ν = acos((dot(r_xyz,e_xyz))/(r*norm(e_xyz)))
  Ω = dot(n_xyz,[0,1,0]) > 0. ? Ω : -Ω
  ω = dot(e_xyz,[0,0,1]) > 0. ? ω : -ω
  ν = dot(r_xyz,v_xyz) > 0. ? ν : -ν
  return [a,e,i,Ω,ω,ν]

end

function total_manuever(r1::Float64,
                        r2::Float64,
                        i1::Float64,
                        i2::Float64,
                        i3::Float64,
                        i4::Float64,
                        μ::Float64)
  if i2 > i3
    return Inf
  else
    a1 = r1
    a2 = (r1+r2)/2
    e1 = 0
    e2 = (r2-r1)/(r2+r1)
    v1_pre = oe_to_xyz([a1, 0., i1, 0., 0., 0.], μ)[4:6]
    v1_post = oe_to_xyz([a2, e2, i2, 0., 0., 0.], μ)[4:6]
    v2_pre = oe_to_xyz([a2, e2, i2, 0., 0., π], μ)[4:6]
    v2_post = oe_to_xyz([a2, e2, i3, 0., 0., π], μ)[4:6]
    v3_pre = oe_to_xyz([a2, e2, i3, 0., 0., 0.], μ)[4:6]
    v3_post = oe_to_xyz([a1, e1, i4, 0., 0., 0.], μ)[4:6]
    Δv1 = norm(v1_post-v1_pre)
    Δv2 = norm(v2_post-v2_pre)
    Δv3 = norm(v3_post-v3_pre)
    return Δv1 + Δv2 + Δv3
  end
end

function plot_maneuver(r1::Float64, r2::Float64, Δ_i::Float64,  μ::Float64, lim::Float64)
  N = 200
  node1_inc = node3_inc = collect(range(0, length=N, stop=lim))
  z = zeros(Float64, N, N)
  base = zeros(Float64, N, N, 2)
  for i in 1:N
    for j in 1:N
      z[i,j] = total_manuever(r1, r2, 0., node1_inc[i], Δ_i-node3_inc[j], Δ_i, μ)
      base[i,j,:] = [node1_inc[i], node3_inc[j]]
    end
  end

  data = contour(;z=z, x=node1_inc, y=node3_inc,
                 name="ΔV", contours=attr(size=0.2),
                 colorbar=attr(title=attr(text="ΔV (km/s)")))
  if lim < Δ_i
    layout = Layout(
      ;title="Total ΔV vs. Δi₁ and Δi₃ for a 90° Bi-Elliptic Transfer (zoomed)",
      width=920,
      height=920,
      xaxis=attr(title=attr(text="Δi₁ (radians)")),
      yaxis=attr(title=attr(text="Δi₃ (radians)")))
    p = Plot(data, layout)
    savefig(p::Union{Plot,PlotlyJS.SyncPlot}, "hw1_4_zoomed.png")
  else
    layout = Layout(
      ;title="Total ΔV vs. Δi₁ and Δi₃ for a 90° Bi-Elliptic Transfer",
      width=920,
      height=920,
      xaxis=attr(title=attr(text="Δi₁ (radians)")),
      yaxis=attr(title=attr(text="Δi₃ (radians)")))
    p = Plot(data, layout)
    savefig(p::Union{Plot,PlotlyJS.SyncPlot}, "hw1_4.png")
  end
  display(plot(p))

  min, position = findmin(z)
  angles = base[position, :]
  println("Found min: ", min)
  println("At: ", 180/π*angles[1], "°, ", 180/π*angles[2], "°")
  while true
    println("Ctrl+C to close")
    sleep(10)
  end
end

end