text:
"""
    Lagrange{N, S, R, T}

Represent a polynomial in Lagrange form using nodes `xs`, weights `ws`, and coefficients `coeffs`.
The Lagrange form does polynomial interpolation between `xs` and `ys` through `p(x) = Σ_{0..n} ℓ_i(x) y_i`,
where if `ℓ(x) = prod(x-x_i)`, `w_i = 1/prod_{j≠i}(x_i - x_j)`, then `ℓ_i(x) = ℓ(x) w_i/(x-x_i)`. The `ℓ_i`
satisfy `ℓ_i(x_j) = δ_{ij}`, so the coefficients are just the `ys`.

```jldoctest Lagrange
julia> using Polynomials, SpecialPolynomials

julia> p =  Lagrange([1,2,3], [1,2,3])
Lagrange(1⋅ℓ^2_0(x) + 2⋅ℓ^2_1(x) + 3⋅ℓ^2_2(x))

julia> p.([1,2,3]) # the coefficients
3-element Array{Int64,1}:
 1
 2
 3

julia> convert(Polynomial,  p)
Polynomial(1.0*x)
```

The instances hold the nodes and weights, which are necessary for
representation, so the type alone can not be used for functions such
as `variable` or `convert(Lagrange, ...)`. For the former we can  use an instance, for the latter we can use `fit`:

```jldoctest Lagrange
julia> p =  Lagrange([1,2,3], [1,2,3])
Lagrange(1⋅ℓ^2_0(x) + 2⋅ℓ^2_1(x) + 3⋅ℓ^2_2(x))

julia> variable(p)
Lagrange(1⋅ℓ^2_0(x) + 2⋅ℓ^2_1(x) + 3⋅ℓ^2_2(x))

julia> q = Polynomial([0,0,1])
Polynomial(x^2)

julia> qq = fit(Lagrange, p.xs, q)
Lagrange(1⋅ℓ^2_0(x) + 4⋅ℓ^2_1(x) + 9⋅ℓ^2_2(x))

julia> convert(Polynomial, qq)
Polynomial(1.0*x^2)
```

Interpolating polynomials suffer from the Runge phenomenon unless the
nodes are well chosen. For `P=Chebyshvev` and `P=ChebyshevU`, the
function `SpecialPolynomials.lagrange_barycentric_nodes_weights(P, n)`
will return a good choice of `n+1` points over `[-1,1]` along with
precomputed weights.

```jldoctest Lagrange
julia> xs, ws = SpecialPolynomials.lagrange_barycentric_nodes_weights(Chebyshev, 64);

julia> f(x) = exp(-x)*sinpi(x)
f (generic function with 1 method)

julia> p = fit(Lagrange, xs, f.(xs));

julia> degree(p)
64

julia> maximum(abs.(f(x) - p(x) for x in range(-1, 1, length=20))) <= 1e-14
true
```

!!! note
    The above example  is  more directly  done through `fit(Chebyshev, f, 64)`, though  the resulting
    polynomial will reference a different  basis.

"""
struct Lagrange{N, S<:Number, R <: Number, T <: Number} <: AbstractInterpolatingPolynomial{T}
    xs::Vector{S}
    ws::Vector{R}
    coeffs::Vector{T}
    var::Symbol
    function  Lagrange(xs::Vector{S}, ws::Vector{R}, coeffs::Vector{T}, var::Symbol=:x) where {S,R,T}
        xs = unique(xs)
        N = length(xs)
        N == length(coeffs) || throw(ArgumentError("the unique xs and the coeffs must have the same length"))
        new{N, S,R,T}(xs, ws, coeffs, var)
    end
    function  Lagrange(xs::Vector{S}, coeffs::Vector{T}, var::Symbol=:x) where {S,T}
        xs = unique(xs)
        N = length(xs)
        ws = lagrange_barycentric_weights(xs)
        R = eltype(ws)
        length(coeffs) == length(xs) || throw(ArgumentError("the unique xs and the coeffs must have the same length"))
        new{N, S, R, T}(xs, ws, coeffs, :x)
    end
end

export Lagrange

basis_symbol(::Type{Lagrange{N,S,R,T}}) where {N,S,R,T} = "ℓ^$(N-1)"

## Boilerplate code reproduced here, as there are three type parameters
Base.convert(::Type{P}, p::P) where {P <: Lagrange} =  p
Base.convert(::Type{Lagrange{N,S,R,T}}, p::Lagrange) where {N,S,R,T} = Lagrange{N,S,R,T}(p.xs, p.ws, coeffs(p), p.var)
Base.promote_rule(::Type{Lagrange{N, S, R, T}}, ::Type{Lagrange{N, S, R, Q}}) where {N, S, R, T, Q <: Number} = Lagrange{N, S, R, promote_type(T,Q)}
Base.promote_rule(::Type{Lagrange{N, S, R, T}}, ::Type{Q}) where {N, S, R, T, Q <: Number} = Lagrange{N, S, R, promote_type(T,Q)}

Polynomials.domain(::Type{<:Lagrange}) = Polynomials.Interval(-Inf, Inf)
function Polynomials.variable(p::Lagrange{S, R, T}, var::Polynomials.SymbolLike=:x) where {S, R, T}
    _fit(Lagrange, p.xs, p.ws, p.xs, var)
end
function Polynomials.one(p::Lagrange)
    _fit(Lagrange, p.xs, p.ws, ones(eltype(p.xs),length(p.xs)), p.var)
end
Polynomials.zero(p::Lagrange{N, S, R, T}) where {N,S,R,T} = 0*p

##  Evaluation
##  From https://people.maths.ox.ac.uk/trefethen/barycentric.pdf
function (p::Lagrange)(x::Number)
    ws, xs, cs = p.ws, p.xs, coeffs(p)
    a = b = zero(x/1)
    for j in eachindex(xs)
        Δ = (x - xs[j])
        iszero(Δ) && return cs[j] # xj ∈ xs, so sum c_j l_i(xj) = sum c_j δ_{ij} = c_j
        l = ws[j]/Δ
        a += l * cs[j]
        b += l
   end
   a/b
end

## convert to poly
## avoids division in barycentric form, which isn't available for Polynomial  arguments, such as `x-xs[j]`
## This has numeric issues, unlike the barycentric form used for evaluation
function  Base.convert(Q::Type{<:Polynomial},  p::Lagrange{N, S,R,T})  where  {N, S, R, T}
    q = zero(Q)/1
    x =  variable(q)
    xs = p.xs
    cs = coeffs(p)
    ws = p.ws

    if length(xs) == 1
       return cs[1] + 0*x #Polynomial(cs[1]*ones(T,1))
    else
        l = fromroots(Polynomial, xs)
        for i in  eachindex(xs)
            li = prod(x - xs[j] for j in eachindex(xs) if j != i);
            q += li * ws[i] * cs[i]
        end
    end
    q
end

## return w⁻¹s = {1/w_j, 0 <= j <= n} as 1 based ws[j+1] = w_j
"""
   lagrange_barycentric_weights(xs)

return [1/w_i for 0 <i <= n] where wi = ∏_{j≠i} (xi - xj)
"""
function lagrange_barycentric_weights(xs)
    n =  length(xs)
    ws =  ones(eltype(xs), n)
    for j in 2:n
        for k in 1:j
            ws[k] =  (xs[k] - xs[j]) * ws[k]
        end
        ws[j] = prod(xs[j]-xs[i] for i in 1:j-1)
    end
    1 ./ ws
end

# https://people.maths.ox.ac.uk/trefethen/barycentric.pdf (5.1)
function lagrange_barycentric_weights(xs::Union{AbstractRange{T},AbstractUnitRange{T}}) where {T}
    n = length(xs)-1
    ws = zeros(float(T), n+1)
    o = one(T)
    for j = 0:n
        ws[1+j] = o * binomial(n, j)
        o *= -1
    end
    ws
end

##  add a new point
## check for  xk ∈ xs?
"""
    update_lagrange_barycentric_weights(ws, xs, xn1)

Adds `xn1` to the nodes `xs` and updates the weights `ws` accordingly.
"""
function update_lagrange_barycentric_weights(ws, xs, xn1)
    ys = [wj/(xj - xn1) for (wj,xj) in zip(ws, xs)]
    wn1 = one(eltype(ws))
    for xj in xs
        wn1 *= (xn1 - xj)
    end
    push!(ys, 1/wn1)
    ys
end

"""
    lagrange_barycentric_nodes_weights(::Type{<:SpecialPolynomial}, n::Int)

Return a collection of `n+1` nodes and weights the given family of
polynomials. Defined for `ChebyshevT` and `ChebyshevU` familes, as there
are `O(n)` algorithms available. Otherwise, `lagrange_barycentric_weights`
is needed.

"""
lagrange_barycentric_nodes_weights(::Type{<:AbstractSpecialPolynomial}, n::Int) = throw(MethodError())

## can add easily if the nodes are shared
## otherwise, this needs to interpolate
function Base.:+(p1::Lagrange{N,S,R,T}, p2::Lagrange{M,S1,R1,T1}) where {N,S,R,T,M,S1,R1,T1}
    p1.var == p2.var || throw(ArgumentError("p1 and p2 must share the same variable"))

    # can just add if using the same set of nodes
    if N == M && p1.xs == p2.xs
        cs =  coeffs(p1) +  coeffs(p2)
        return Lagrange(p1.xs, p1.ws, cs, p1.var)
    else
        p,q = N >= M ? (p1, p2) : (p2, p1)
        xs, ws, cs = p.xs, p.ws, copy(coeffs(p))
        for i in eachindex(xs)
            cs[i] += q(xs[i])
        end
        return _fit(Lagrange, xs,  ws, cs, p1.var)
    end
end

## XXX is there a better way?
## This refits using  updated weights,
## but this can be wildly inaccurate
function Base.:*(p1::Lagrange{N,S,R,T}, p2::Lagrange{M,S1,R1,T1}) where {N,S,R,T,M,S1,R1,T1}

    p1.var == p2.var || throw(ArgumentError("p1 and p2 must share the same variable"))

    p, q =  N >= M ? (p1, p2)  : (p2, p1)
    S2 = typeof(one(S) * one(S1)/1)
    xs = convert(Vector{S2}, copy(p.xs))
    ws = copy(p.ws)

    # new_xs ## XXX this is where the work is
    new_xs = _new_nodes(xs, q.xs)
    for y in new_xs
        ws = update_lagrange_barycentric_weights(ws, xs, y)
        push!(xs, y)
    end
    ys = p.(xs) .* q.(xs)
    _fit(Lagrange, xs,  ws, ys, p1.var)

end

# handle scalar case
function Base.:+(p::Lagrange, c::Number)
    q = Lagrange(p.xs[1:1], [c], p.var)
    p + q
end

Base.:*(c::Number, p::Lagrange) = p*c
function Base.:*(p::P, c::Number) where {P <: Lagrange}
    Lagrange(p.xs, c*coeffs(p), p.var)
end

Base.:-(p::P) where {P<:Lagrange} = (-1)*p

## short cut if weights are known.
_fit(P::Type{Lagrange}, xs, ws, ys, var=:x) = Lagrange(xs, ws, ys, var)

function Polynomials.fit(P::Type{<:Lagrange},
                         xs::AbstractVector{S},
                         ys::AbstractVector{T};
                         var = :x) where {S, T}
    ws = lagrange_barycentric_weights(xs)
    _fit(P, xs, ws, ys, var)
end

function Polynomials.fit(P::Type{<:Lagrange}, xs::AbstractVector{S}, f; var=:x) where {S}
     ws = lagrange_barycentric_weights(xs)
    _fit(P, xs, ws, f.(xs), var)
end


text:
module mocaputil

using StatsBase, Statistics, ArgCheck
import Random: shuffle

export transform, scale_transform
export fit
export invert
export MyStandardScaler

"""
    fit(MyStandardScaler, X, dims)

Fit a standardisation to a matrix `X` s.t. that the rows/columns have mean zero and standard deviation 1.
This operation calculates the mean and standard deviation and outputs a MyStandardScaler object `s` which
allows this same standardisation to be fit to any matrix using `transform(s, Y)` for some matrix `Y`. Note
that the input matrix `X` is *not* transformed by this operation. Instead use the above `transform` syntax
on `X`.

Note a couple of addendums:

1. Any columns/rows with constant values will result in a standard deviation of 1.0, not 0. This is to avoid NaN errors from the transformation (and it is a natural choice).
2. If only a subset of the rows/columns should be standardised, an additional argument of the indices may be given as:

    `fit(MyStandardScaler, X, operate_on, dims)`

    The subset to operate upon will be maintained through all `transform` and `invert` operations.
"""
mutable struct MyStandardScaler{T}
    μ::Array{T,1}
    σ::Array{T,1}
    operate_on::Array{L where L <: Int,1}
    dims::Int
end

Base.copy(s::MyStandardScaler) = MyStandardScaler(copy(s.μ), copy(s.σ), copy(s.operate_on), copy(s.dims))

StatsBase.fit(::Type{MyStandardScaler}, X::Matrix, dims::Int=1) = MyStandardScaler(vec(mean(X, dims=dims)),
    vec(std(X, dims=dims)), collect(1:size(X,3-dims)), dims) |> post_fit

StatsBase.fit(::Type{MyStandardScaler}, X::Matrix, operate_on::Array{L where L <: Int}, dims::Int=1) =
    MyStandardScaler(vec(mean(X, dims=dims))[operate_on],
    vec(std(X, dims=dims))[operate_on],
    operate_on, dims) |> post_fit

function post_fit(s::MyStandardScaler)
    bad_ixs = (s.σ .== 0)
    s.σ[bad_ixs] .= 1
    s
end

function scale_transform(s::MyStandardScaler, X::Matrix, dims::Int=s.dims)
    (dims != s.dims) && @warn "dim specified in transform is different to specification during fit."
    tr = dims == 1 ? transpose : identity
    tr2 = dims == 2 ? transpose : identity
    if s.operate_on == 1:size(X, 3-dims)
        out = (X .- tr(s.μ)) ./ tr(s.σ)
    else
        out = tr2(copy(X))
        s_sub = copy(s)
        s_sub.operate_on = 1:length(s.operate_on)
        s_sub.dims = 1
        out[:, s.operate_on] = transform(s_sub, out[:, s.operate_on], 1)
        out = tr2(out)
    end
    out
end

function invert(s::MyStandardScaler, X::Matrix, dims::Int=s.dims)
    (dims != s.dims) && @warn "dim specified in inversion is different to specification during fit."
    tr = dims == 1 ? transpose : identity
    tr2 = dims == 2 ? transpose : identity
    if s.operate_on == 1:size(X, 3-dims)
        out = (X .* tr(s.σ)) .+ tr(s.μ)
    else
        out = tr2(copy(X))
        s_sub = copy(s)
        s_sub.operate_on = 1:length(s.operate_on)
        s_sub.dims = 1
        out[:, s.operate_on] = invert(s_sub, out[:, s.operate_on], 1)
        out = tr2(out)
    end
    out
end

"""
DataIterator is conceptually similar to a DataLoader class in PyTorch. It has the following methods:

* `iterate`   (obtain the next iteration of data)
* `weights`   (extract all weights from the DataIterator for each datapoint and normalize)
* `copy`
* `shuffle`   (shuffles the order of the data)
* `indexed_shuffle`  (shuffles the order of the data, and additionally spits out the indices of the new order)
"""
struct DataIterator
    data::Array{D}  where {D <: Dict}
    batch_size::Int
    min_size::Int
    start::Int
    DataIterator(d, b_sz, m_sz, start) = begin; @assert m_sz < b_sz; new(d, b_sz, m_sz, start); end
end
DataIterator(data, batch_size; min_size=1, start=0) = DataIterator(data, batch_size, min_size, start)
Base.copy(di::DataIterator) = DataIterator(deepcopy(di.data), copy(di.batch_size), copy(di.min_size), copy(di.start))

function Base.iterate(iter::DataIterator, state=(1, iter.start))
    element, ix = state

    (element > length(iter.data)) && return nothing
    while ix + iter.min_size > size(iter.data[element][:Y], 2)
        element += 1
        ix = iter.start
        (element > length(iter.data)) && return nothing
    end
    new_state = ix == iter.start   # not in while, since 1st iterate in general won't use this loop.

    chunk  = iter.data[element]
    cur_length = size(chunk[:Y], 2)
    eix  = min(ix + iter.batch_size, cur_length)
    ix += 1

    return ((chunk[:Y][:,ix:eix], chunk[:U][:,ix:eix], new_state), (element, eix))
end

function weights(iter::DataIterator; as_pct::Bool=true)
    w = [size(y,2) for (y, u, new_state) in iter]
    return as_pct ? w / sum(w) : w
end

function Base.length(iter::DataIterator)
    map(iter.data) do x
        d, r = divrem(size(x[:Y],2)-iter.start, iter.batch_size);
        d + (r >= iter.min_size);
        end |> sum
end

function shuffle(di::DataIterator)
    newdata = shuffle([Dict(:Y=>cY, :U=>cU) for (cY, cU, h0) in di])
    DataIterator(newdata, di.batch_size, 1, di.start)
end

function indexed_shuffle(di::DataIterator)
    order = shuffle(1:length(di))
    dataArray = [Dict(:Y=>cY, :U=>cU) for (cY, cU, h0) in di]
    newdata = [dataArray[i] for i in order]
    order, DataIterator(newdata, di.batch_size, 1, di.start)
end

function get_file_pos_from_iter(di::DataIterator, styles_lkp::Vector,
        style_ix::Int, ix::Int)
    ffs = vcat([s for (i,s) in enumerate(styles_lkp) if i != style_ix]...)
    e, file_offset = 1, di.start
    for i in 1:ix-1
        _, (e, file_offset) = iterate(di, (e, file_offset))
    end
    return ffs[e], file_offset+1
end

function get_file_pos_from_iter_test(di::DataIterator, styles_lkp::Vector,
        style_ix::Int, ix::Int)
    get_file_pos_from_iter(di, [styles_lkp[style_ix]], 0,ix)
end

end

ext: jl

max_stars_repo_name: ornithos/pytorch-mtds-mocap

max_issues_repo_name: ornithos/pytorch-mtds-mocap

max_forks_repo_name: ornithos/pytorch-mtds-mocap


text:
include("Cosmology.jl")
using Cosmology
using Base.Test

function test_approx_eq_rtol(va, vb, rtol, astr, bstr)
    diff = maximum(abs(va - vb))
    tol = rtol*max(maximum(abs(va)), maximum(abs(vb)))
    if diff > tol
        sdiff = string("|", astr, " - ", bstr, "| <= ", tol)
        error("assertion failed: ", sdiff,
              "\n  ", astr, " = ", va,
              "\n  ", bstr, " = ", vb,
              "\n  difference = ", diff, " > ", tol)
    end
end

macro test_approx_eq_rtol(a, b, c)
    :(test_approx_eq_rtol($(esc(a)), $(esc(b)), $(esc(c)), $(string(a)), $(string(b))))
end

# values from http://icosmos.co.uk/

dist_rtol = 1e-6
age_rtol = 2e-4

c = cosmology(h=0.7, OmegaM=0.3, OmegaR=0)
@test_approx_eq_rtol angular_diameter_dist_mpc(c,1) 1651.9145 dist_rtol
@test_approx_eq_rtol comoving_radial_dist_mpc(c,1) 3303.829 dist_rtol
@test_approx_eq_rtol comoving_volume_gpc3(c,1) 151.0571 dist_rtol
@test_approx_eq_rtol luminosity_dist_mpc(c,1) 6607.6579 dist_rtol
@test_approx_eq_rtol age_gyr(c,0) 13.4694 age_rtol
@test_approx_eq_rtol age_gyr(c,1) 5.7527 age_rtol
@test_approx_eq_rtol lookback_time_gyr(c,1) 13.4694-5.7527 age_rtol

c = cosmology(h=0.7, OmegaK=0.1, OmegaM=0.3, OmegaR=0)
@test_approx_eq_rtol angular_diameter_dist_mpc(c,1) 1619.9588 dist_rtol
@test_approx_eq_rtol comoving_radial_dist_mpc(c,1) 3209.784 dist_rtol
@test_approx_eq_rtol comoving_volume_gpc3(c,1) 140.0856 dist_rtol
@test_approx_eq_rtol luminosity_dist_mpc(c,1) 6479.8352 dist_rtol
@test_approx_eq_rtol age_gyr(c,0) 13.064 age_rtol
@test_approx_eq_rtol age_gyr(c,1) 5.5466 age_rtol
@test_approx_eq_rtol lookback_time_gyr(c,1) 13.064-5.5466 age_rtol

c = cosmology(h=0.7, OmegaK=-0.1, OmegaM=0.3, OmegaR=0)
@test_approx_eq_rtol angular_diameter_dist_mpc(c,1) 1686.5272 dist_rtol
@test_approx_eq_rtol comoving_radial_dist_mpc(c,1) 3408.937 dist_rtol
@test_approx_eq_rtol comoving_volume_gpc3(c,1) 163.8479 dist_rtol
@test_approx_eq_rtol luminosity_dist_mpc(c,1) 6746.1088 dist_rtol
@test_approx_eq_rtol age_gyr(c,0) 13.925 age_rtol
@test_approx_eq_rtol age_gyr(c,1) 5.9868 age_rtol
@test_approx_eq_rtol lookback_time_gyr(c,1) 13.925-5.9868 age_rtol

c = cosmology(h=0.7, OmegaM=0.3, OmegaR=0, w0=-0.9, wa=0.1)
@test_approx_eq_rtol angular_diameter_dist_mpc(c,1) 1612.0585 dist_rtol
@test_approx_eq_rtol comoving_radial_dist_mpc(c,1) 3224.1169 dist_rtol
@test_approx_eq_rtol comoving_volume_gpc3(c,1) 140.3851 dist_rtol
@test_approx_eq_rtol luminosity_dist_mpc(c,1) 6448.2338 dist_rtol
@test_approx_eq_rtol age_gyr(c,0) 13.1915 age_rtol
@test_approx_eq_rtol age_gyr(c,1) 5.6464 age_rtol
@test_approx_eq_rtol lookback_time_gyr(c,1) 13.1915-5.6464 age_rtol

c = cosmology(h=0.7, OmegaK=0.1, OmegaM=0.3, OmegaR=0, w0=-0.9, wa=0.1)
@test_approx_eq_rtol angular_diameter_dist_mpc(c,1) 1588.0181 dist_rtol
@test_approx_eq_rtol comoving_radial_dist_mpc(c,1) 3147.6227 dist_rtol
@test_approx_eq_rtol comoving_volume_gpc3(c,1) 132.0466 dist_rtol
@test_approx_eq_rtol luminosity_dist_mpc(c,1) 6352.0723 dist_rtol
@test_approx_eq_rtol age_gyr(c,0) 12.8488 age_rtol
@test_approx_eq_rtol age_gyr(c,1) 5.4659 age_rtol
@test_approx_eq_rtol lookback_time_gyr(c,1) 12.8488-5.4659 age_rtol

c = cosmology(h=0.7, OmegaK=-0.1, OmegaM=0.3, OmegaR=0, w0=-0.9, wa=0.1)
@test_approx_eq_rtol angular_diameter_dist_mpc(c,1) 1637.5993 dist_rtol
@test_approx_eq_rtol comoving_radial_dist_mpc(c,1) 3307.9932 dist_rtol
@test_approx_eq_rtol comoving_volume_gpc3(c,1) 149.8301 dist_rtol
@test_approx_eq_rtol luminosity_dist_mpc(c,1) 6550.3973 dist_rtol
@test_approx_eq_rtol age_gyr(c,0) 13.5702 age_rtol
@test_approx_eq_rtol age_gyr(c,1) 5.8482 age_rtol
@test_approx_eq_rtol lookback_time_gyr(c,1) 13.5702-5.8482 age_rtol


text:
# Sparse Matrix COO
struct SparseMatrixCOO{Tv, Ti}
    m::Int
    n::Int
    i::Vector{Ti}
    j::Vector{Ti}
    v::Vector{Tv}
end

to_sparse_csc(coo::SparseMatrixCOO) = sparse(coo.i, coo.j, coo.v, coo.m, coo.n)

to_full(coo::SparseMatrixCOO) = full(to_sparse_csc(coo))

to_python(coo::SparseMatrixCOO) = scipy_sps["coo_matrix"](
    (coo.v, (coo.i-1, coo.j-1)), shape=(coo.m, coo.n))

# Linear Operator
struct LinearOperator
    m::Int
    n::Int
    matvec::Function
end

function Base.dot(op::LinearOperator, x::Vector)
    return op.matvec(x)
end

to_python(op::LinearOperator) = scipy_spslin["LinearOperator"]((op.m, op.n),
                                                               matvec=op.matvec)


text:
module TestFinSets
using Test

using Catlab.Theories, Catlab.CategoricalAlgebra
using Catlab.CategoricalAlgebra.FinSets

# Category of finite sets: skeleton
###################################

f = FinFunction([1,3,4], 5)
g = FinFunction([1,1,2,2,3], 3)
h = FinFunction([3,1,2], 3)

# Evaluation.
rot3(x) = (x % 3) + 1
@test map(f, 1:3) == [1,3,4]
@test map(FinFunction(rot3, 3, 3), 1:3) == [2,3,1]
@test map(id(FinSet(3)), 1:3) == [1,2,3]

# Pretty-print.
@test sprint(show, FinSet(3)) == "FinSet(3)"
@test sprint(show, f) == "FinFunction([1, 3, 4], 3, 5)"
@test sprint(show, FinFunction(rot3, 3, 3)) ==
  "FinFunction(rot3, FinSet(3), FinSet(3))"
@test sprint(show, id(FinSet(3))) == "FinFunction(identity, FinSet(3))"

# Domains and codomains.
@test dom(f) == FinSet(3)
@test codom(f) == FinSet(5)
@test dom(id(FinSet(3))) == FinSet(3)
@test codom(id(FinSet(3))) == FinSet(3)

# Composition and identities.
@test compose(f,g) == FinFunction([1,2,2], 3)
@test compose(g,h) == FinFunction([3,3,1,1,2], 3)
@test compose(compose(f,g),h) == compose(f,compose(g,h))
@test compose(id(dom(f)), f) == f
@test compose(f, id(codom(f))) == f

# Limits
#-------

# Terminal object.
@test ob(terminal(FinSet{Int})) == FinSet(1)
@test delete(terminal(FinSet{Int}), FinSet(3)) == FinFunction([1,1,1])

# Binary product.
lim = product(FinSet(2), FinSet(3))
@test ob(lim) == FinSet(6)
@test force(proj1(lim)) == FinFunction([1,2,1,2,1,2])
@test force(proj2(lim)) == FinFunction([1,1,2,2,3,3])

lim = product(FinSet(4), FinSet(3))
f, g = FinFunction([2,1,4]), FinFunction([1,3,2])
@test force(pair(lim,f,g) ⋅ proj1(lim)) == f
@test force(pair(lim,f,g) ⋅ proj2(lim)) == g

# N-ary product.
lim = product([FinSet(2), FinSet(3)])
@test ob(lim) == FinSet(6)
@test force.(legs(lim)) ==
  [FinFunction([1,2,1,2,1,2]), FinFunction([1,1,2,2,3,3])]
@test ob(product(FinSet{Int}[])) == FinSet(1)

lim = product([FinSet(4), FinSet(3)])
@test force(pair(lim,[f,g]) ⋅ first(legs(lim))) == f
@test force(pair(lim,[f,g]) ⋅ last(legs(lim))) == g

# Equalizer.
f, g = FinFunction([1,2,4,3]), FinFunction([3,2,4,1])
eq = equalizer(f,g)
@test incl(eq) == FinFunction([2,3], 4)
@test incl(equalizer([f,g])) == incl(eq)
@test factorize(eq, FinFunction([2,3,2])) == FinFunction([1,2,1])

# Equalizer in case of identical functions.
f = FinFunction([4,2,3,1], 5)
eq = equalizer(f,f)
@test incl(eq) == force(id(FinSet(4)))
@test incl(equalizer([f,f])) == incl(eq)
@test factorize(eq, FinFunction([2,1,3,3])) == FinFunction([2,1,3,3], 4)

# Equalizer matching nothing.
f, g = id(FinSet(5)), FinFunction([2,3,4,5,1], 5)
eq = equalizer(f,g)
@test incl(eq) == FinFunction(Int[], 5)
@test incl(equalizer([f,g])) == incl(eq)
@test factorize(eq, FinFunction(Int[])) == FinFunction(Int[])

# Pullback.
lim = pullback(FinFunction([1,1,3,2],4), FinFunction([1,1,4,2],4))
@test ob(lim) == FinSet(5)
@test force(proj1(lim)) == FinFunction([1,2,1,2,4], 4)
@test force(proj2(lim)) == FinFunction([1,1,2,2,4], 4)

f, g = FinFunction([1,2,4]), FinFunction([2,1,4])
@test force(pair(lim,f,g) ⋅ proj1(lim)) == f
@test force(pair(lim,f,g) ⋅ proj2(lim)) == g

# Pullback from a singleton set: the degenerate case of a product.
lim = pullback(FinFunction([1,1]), FinFunction([1,1,1]))
@test ob(lim) == FinSet(6)
@test force(proj1(lim)) == FinFunction([1,2,1,2,1,2])
@test force(proj2(lim)) == FinFunction([1,1,2,2,3,3])

f, g = FinFunction([1,1,2]), FinFunction([3,2,1])
@test force(pair(lim,f,g) ⋅ proj1(lim)) == f
@test force(pair(lim,f,g) ⋅ proj2(lim)) == g

# Pullback using generic limit interface
f, g = FinFunction([1,1,3,2],4), FinFunction([1,1,4,2],4)
lim = limit(FreeDiagram([FinSet(4),FinSet(4),FinSet(4)], [(1,3,f),(2,3,g)]))
@test ob(lim) == FinSet(5)
@test force.(legs(lim)[1:2]) ==
  [FinFunction([1,2,1,2,4],4), FinFunction([1,1,2,2,4],4)]

# Colimits
#---------

# Initial object.
@test ob(initial(FinSet{Int})) == FinSet(0)
@test create(initial(FinSet{Int}), FinSet(3)) == FinFunction(Int[], 3)

# Binary coproduct.
colim = coproduct(FinSet(2), FinSet(3))
@test ob(colim) == FinSet(5)
@test coproj1(colim) == FinFunction([1,2], 5)
@test coproj2(colim) == FinFunction([3,4,5], 5)

f, g = FinFunction([3,5], 5), FinFunction([1,2,3], 5)
@test force(coproj1(colim) ⋅ copair(colim,f,g)) == f
@test force(coproj2(colim) ⋅ copair(colim,f,g)) == g

# N-ary coproduct.
colim = coproduct([FinSet(2), FinSet(3)])
@test ob(colim) == FinSet(5)
@test legs(colim) == [FinFunction([1,2], 5), FinFunction([3,4,5], 5)]
@test ob(coproduct(FinSet{Int}[])) == FinSet(0)

@test force(first(legs(colim)) ⋅ copair(colim,[f,g])) == f
@test force(last(legs(colim)) ⋅ copair(colim,[f,g])) == g

# Coequalizer from a singleton set.
f, g = FinFunction([1], 3), FinFunction([3], 3)
coeq = coequalizer(f,g)
@test proj(coeq) == FinFunction([1,2,1], 2)
@test proj(coequalizer([f,g])) == proj(coeq)
@test factorize(coeq, FinFunction([4,1,4])) == FinFunction([4,1])
@test_throws AssertionError factorize(coeq, FinFunction([3,1,4]))

# Coequalizer in case of identical functions.
f = FinFunction([4,2,3,1], 5)
coeq = coequalizer(f,f)
@test proj(coeq) == force(id(FinSet(5)))
@test proj(coequalizer([f,f])) == proj(coeq)
@test factorize(coeq, FinFunction([2,1,3,3,4])) == FinFunction([2,1,3,3,4])

# Coequalizer identifying everything.
f, g = id(FinSet(5)), FinFunction([2,3,4,5,1], 5)
coeq = coequalizer(f,g)
@test proj(coeq) == FinFunction(fill(1,5))
@test proj(coequalizer([f,g])) == proj(coeq)
@test factorize(coeq, FinFunction(fill(3,5))) == FinFunction([3])

# Pushout from the empty set: the degenerate case of the coproduct.
f, g = FinFunction(Int[], 2), FinFunction(Int[], 3)
colim = pushout(f,g)
@test ob(colim) == FinSet(5)
@test coproj1(colim) == FinFunction([1,2], 5)
@test coproj2(colim) == FinFunction([3,4,5], 5)

h, k = FinFunction([3,5], 5), FinFunction([1,2,3], 5)
@test force(coproj1(colim) ⋅ copair(colim,h,k)) == h
@test force(coproj2(colim) ⋅ copair(colim,h,k)) == k

# Pushout from a singleton set.
f, g = FinFunction([1], 2), FinFunction([2], 3)
colim = pushout(f,g)
@test ob(colim) == FinSet(4)
ι1, ι2 = colim
@test compose(f,ι1) == compose(g,ι2)
@test ι1 == FinFunction([1,2], 4)
@test ι2 == FinFunction([3,1,4], 4)

h, k = FinFunction([3,5]), FinFunction([1,3,5])
@test force(coproj1(colim) ⋅ copair(colim,h,k)) == h
@test force(coproj2(colim) ⋅ copair(colim,h,k)) == k
k = FinFunction([1,2,5])
@test_throws AssertionError copair(colim,h,k)

# Same thing with generic colimit interface
diag = FreeDiagram([FinSet(1),FinSet(2),FinSet(3)],[(1,2,f), (1,3,g)])
colim = colimit(diag)
@test ob(colim) == FinSet(4)
_, ι1, ι2 = colim
@test compose(f,ι1) == compose(g,ι2)
@test ι1 == FinFunction([1,2], 4)
@test ι2 == FinFunction([3,1,4], 4)

# Pushout from a two-element set, with non-injective legs.
f, g = FinFunction([1,1], 2), FinFunction([1,2], 2)
colim = pushout(f,g)
@test ob(colim) == FinSet(2)
ι1, ι2 = colim
@test compose(f,ι1) == compose(g,ι2)
@test ι1 == FinFunction([1,2], 2)
@test ι2 == FinFunction([1,1], 2)

# Same thing with generic colimit interface
diag = FreeDiagram([FinSet(2),FinSet(2),FinSet(2)],[(1,2,f),(1,3,g)])
colim = colimit(diag)
@test ob(colim) == FinSet(2)
_, ι1, ι2 = colim
@test compose(f,ι1) == compose(g,ι2)
@test ι1 == FinFunction([1,2], 2)
@test ι2 == FinFunction([1,1], 2)

end


text:
module TrainingTests

using Photon, Test
using Photon.Layers
using Photon.Losses
using Photon.Callbacks: SilentMeter

function simple_conv_model()
    model = Sequential(
        Conv2D(16, 3, :relu),
        Conv2D(32, 3, :relu),
        MaxPool2D(),
        Dense(32, :relu),
        Dense(10)
    )
    return model
end

function test_save_load()
    m = simple_conv_model()
    X = KorA(randn(Float32,28,28,1,4))
    pred = m(X)
    workout = Workout(m, MSELoss())
    f = saveworkout(workout)
    workout2 = loadworkout(f)
    rm(f, force=true)
    pred2 = workout2.model(X)
    @assert pred == pred2
end

function test_predict()
    m = simple_conv_model()
    X = KorA(randn(Float32,28,28,1,4))
    workout = Workout(m, MSELoss())
    p1 = predict(workout, X, makebatch=false)
    p2 = predict(workout, X[:,:,:,1])

    # @assert p1[:,:,:,1] == p2
end

function getdata(s=28)
    [(
        KorA(randn(Float32,s,s,1,16)),
        KorA(randn(Float32,10,16))
    ) for i=1:10]
end

function test_train()
    model = simple_conv_model()
    workout = Workout(model, MSELoss())

    data = getdata()
    train!(workout, data, epochs=2, cb=SilentMeter())

    @test workout.epochs == 2
    @test workout.steps == (2 * length(data))
    @test hasmetric(workout, :loss)
end

function test_train_valid()
    model = simple_conv_model()
    workout = Workout(model, MSELoss())

    data_tr = getdata()
    data_val = getdata()
    train!(workout, data_tr, data_val, epochs=2, cb=SilentMeter())

    @test hasmetric(workout, :loss)
    @test hasmetric(workout, :val_loss)
end

function test_validate()
    model = simple_conv_model()
    workout = Workout(model, MSELoss())
    e, s = workout.epochs, workout.steps

    data = getdata()
    validate(workout, first(data))

    @test workout.epochs == e
    @test workout.steps == s
    @test hasmetric(workout, :val_loss)
end

function test_gradients()
    model = simple_conv_model()
    workout = Workout(model, MSELoss())
    e, s = workout.epochs, workout.steps

    data = getdata()
    g = gradients(workout, first(data))

    @test workout.epochs == e
    @test workout.steps == s
end

function test_channel()
    model = simple_conv_model()
    workout = Workout(model, MSELoss())

    minibatch = (KorA(randn(Float32,30,30,3,4)), KorA(randn(Float32,10,4)))

    function randomdata()
        Channel() do channel
            for i in 1:4
                put!(channel, minibatch)
            end
        end
    end

    train!(workout, randomdata, epochs=2,cb=SilentMeter())
    @test workout.epochs == 2
    @test hasmetric(workout, :loss)
end

@testset "Training" begin
    test_train()
    test_train_valid()
    test_validate()
    test_predict()
    test_gradients()
    test_channel()
    test_save_load()
end

end


text:
import Ghost: Tape, Input, Call, mkcall, V

"""
    rename_keys(dict::Dict, subs::Dict)

Create a copy of `dict`, replacing its keys according to mapping `subs`
"""
rename_keys(dct::Dict, subs::Dict) = Dict(get(subs, k, k) => v for (k, v) in pairs(dct))

"""
    unpacked_vars(tape::Ghost.Tape, op::Ghost.Call)

For a multi-output call, find variables on the tape referring to elements of the output.
Example:

    import Ghost: Tape, Input, mkcall

    make_tuple(x) = (x, x + 1)

    tape = Tape()
    x = push!(tape, Input(1.0))
    out = push!(tape, mkcall(make_tuple, x))
    y1 = push!(tape, mkcall(getfield, out, 1))
    y2 = push!(tape, mkcall(getfield, out, 2))

    @assert unpacked_vars(tape, tape[out]) == [y1, y2]

"""
function unpacked_vars(op::Call)
    @assert op.val isa Tuple "Can't unpack non-tuple output"
    tape = op.tape
    # out = V(op)
    vars = Any[nothing for _=1:length(op.val)]
    found = 0
    for id in op.id + 1:length(tape)
        cur = tape[V(id)]
        if (cur isa Call && cur.fn in (getfield, getindex) &&   # is getfield call
            cur.args[1] isa V && cur.args[1].id == op.id)       # with V(op) as the first argument
            # found unpacked var
            idx = cur.args[2]
            idx = idx isa V ? tape[idx].val : idx
            vars[idx] = V(cur)
            # in most cases unpacked vars go right after the multi-ouput,
            # no need to search further
            found += 1
            found == length(vars) && break
        end
    end
    return vars
end


text:
using LinearAlgebra
using Statistics

function computegeometriccenter(points)
    x = [points[i][1] for i in 1:length(points)]
    y = [points[i][2] for i in 1:length(points)]
    z = [points[i][3] for i in 1:length(points)]

    return [mean(x), mean(y), mean(z)]
end

function computesurfacenormal(vectors)
    v1 = vectors[2] - vectors[1]
    v2 = vectors[3] - vectors[1]
    return 0.5*cross(v1, v2)
end

function computemagnitude(vector)
    return sqrt(dot(vector, vector))
end


text:
include("evaluate_lambda.jl")

using PyPlot

α = [-0.5;0.5]
L = round.(Int, exp10.(range(0, stop=4, length=50)))
λlap = [-ℓ*(ℓ+1) for ℓ in L]

for i = 1:length(α)
    δ1 = 0.01
    λ1 = evaluate_lambda(L[end]+1, α[i], δ1)[L.+1]
    δ2 = 0.1
    λ2 = evaluate_lambda(L[end]+1, α[i], δ2)[L.+1]
    δ3 = 1.0
    λ3 = evaluate_lambda(L[end]+1, α[i], δ3)[L.+1]
    δ4 = 2.0
    λ4 = evaluate_lambda(L[end]+1, α[i], δ4)[L.+1]

    clf()
    loglog(L, -λlap, "-k", L, -λ1, ".k", L, -λ2, "xk", L, -λ3, "+k", L, -λ4, "--k")
    xlabel("\$\\ell\$"); ylabel("Numerical Evaluation of \$\\lambda_\\delta(\\ell)\$ at \$\\alpha=$(α[i])\$"); grid()
    legend(["\$-\\lambda_0(\\ell)\$","\$-\\lambda_{$(δ1)}(\\ell)\$","\$-\\lambda_{$(δ2)}(\\ell)\$","\$-\\lambda_{$(δ3 == Int(δ3) ? Int(δ3) : δ3)}(\\ell)\$","\$-\\lambda_{$(δ4 == Int(δ4) ? Int(δ4) : δ4)}(\\ell)\$"])
    savefig("plots/plot_spectrum$(i).pdf")
end


text:
type Game
    board::Board
    player::Player
    level::Int
    previous_levels_score::Int
    players_turn::Bool

    Game(p::Player) = new(Board(),p, 1, 0, true)
end

function next_level!(g::Game)
    g.previous_levels_score += get_score(g.board)
    g.level += 1
    g.board = Board(g.board.height, g.board.width, g.level < 4 ? 10*g.level : 40)
    g.players_turn = true
end

score(g::Game) = g.previous_levels_score + get_score(g.board)
is_over(g::Game) = !g.board.sprite.is_alive
end_turn(g::Game) = g.players_turn = !g.players_turn

function print_score(g::Game, w::Ptr{Void})
    pos = (g.board.height + 3, 1)
    easy_print(w,pos...,string("Score: ",score(g)))
end

function print_level(g::Game, w::Ptr{Void})
    pos = (g.board.height + 4, 1)
    easy_print(w,pos...,string("Level: ",g.level))
end

function print_and_refresh(g::Game,w::Ptr{Void})
    TermWin.erase()
    print_frame(g.board,w)
    print_field(g.board,w)
    print_score(g,w)
    print_level(g,w)
    TermWin.refresh()
end

function play(g::Game, p::Human ,w::Ptr{Void})
    while !is_over(g)
        if !g.board.wait_mode
            print_and_refresh(g,w)
            m = get_valid_move(g.player,g.board, w)
            com = key_to_command(m)
            execute_command(com,g.board)
        end
        process_robot_turn!(g.board)
        scrap_sprite!(g.board)
        switch_active_board!(g.board)
        if g.board.live_robots == 0
            next_level!(g)
        end
    end
end

max_stars_repo_name: RatanRSur/robots-from-scratch

max_issues_repo_name: RatanRSur/robots-from-scratch

max_forks_repo_name: RatanRSur/robots-from-scratch


text:
using Test

using MIToS.Utils
using MIToS.MSA
using MIToS.Information
using MIToS.PDB
using MIToS.SIFTS
using MIToS.Pfam
using LinearAlgebra
using Random
using Statistics            # mean
using DelimitedFiles        # readdlm
using ROCAnalysis           # AUC
using Clustering            # test/MSA/Hobohm.jl
using NamedArrays           # array
using StatsBase             # WeightVec
using PairwiseListMatrices  # getlist

const DATA = "data"

# Utils
@testset "Utils" begin
    include("Utils/GeneralUtils.jl")
end

# MSA
@testset "MSA" begin
    include("MSA/Residues.jl")
    include("MSA/Alphabet.jl")
    include("MSA/ThreeLetters.jl")
    include("MSA/Annotations.jl")
    include("MSA/MultipleSequenceAlignment.jl")
    include("MSA/GeneralParserMethods.jl")
    include("MSA/IO.jl")
    include("MSA/General.jl")
    include("MSA/MSAEditing.jl")
    include("MSA/MSAStats.jl")
    include("MSA/Shuffle.jl")
    include("MSA/Identity.jl")
    include("MSA/Hobohm.jl")
    include("MSA/MSAAnnotations.jl")
    include("MSA/GetIndex.jl")
    include("MSA/Concatenation.jl")
end

# Information
@testset "Information" begin
    include("Information/ContingencyTables.jl")
    include("Information/Counters.jl")
    include("Information/InformationMeasures.jl")
    include("Information/Iterations.jl")
    include("Information/CorrectedMutualInformation.jl")
    include("Information/Gaps.jl")
    include("Information/Externals.jl")
end

# PDB
@testset "PDB" begin
    include("PDB/PDB.jl")
    include("PDB/Contacts.jl")
    include("PDB/Kabsch.jl")
    include("PDB/Internals.jl")
end

# SIFTS
@testset "SIFTS" begin
    include("SIFTS/SIFTS.jl")
end

# Pfam
@testset "Pfam" begin
    include("Pfam/Pfam.jl")
end

# Scripts
@testset "Scripts" begin
    include("Scripts/Template.jl")
    include("Scripts/Scripts.jl")
end

print("""

----- =D -----

""")


text:
using Base.Test
using Snopt
using ForwardDiff
include("dircol.jl")

function gen_usrfun(model, obj, dt, pack, method; grads=:none)
    n,m,N = pack

    # Weights (Trapezoidal)
    weights = get_weights(method,N)

    # Jacobian for Dynamics
    f_aug! = f_augmented!(model!.f, model.n, model.m)
    zdot = zeros(n)
    F(z) = ForwardDiff.jacobian(f_aug!,zdot,z)

    function eval_f(Z)
        X,U = unpackZ(Z,pack)
        cost(X,U,weights,obj)
    end

    function eval_g(Z)
        X,U = unpackZ(Z,pack)
        calc_constraints(X,U,method,dt,model.f)
    end

    function usrfun(Z)
        # Objective Function (Cost)
        J = eval_f(Z)

        # Constraints (dynamics only for now)
        c = Float64[] # No inequality constraints for now
        ceq = eval_g(Z)

        fail = false

        if grads == :none
            return J, c, ceq, fail
        else
            X,U = unpackZ(Z,pack)

            # Gradient of Objective
            if grads == :auto
                grad_f = ForwardDiff.gradient(eval_f,Z)
            else
                grad_f = cost_gradient(X,U,weights,obj)
            end

            # Constraint Jacobian
            jacob_c = Float64[]
            if grads == :auto
                jacob_ceq = ForwardDiff.jacobian(eval_g,Z)
            else
                jacob_ceq = constraint_jacobian(X,U,dt,method,F)
            end

            return J, c, ceq, grad_f, jacob_c, jacob_ceq, fail
        end
    end

    return usrfun, eval_f, eval_g
end

function dircol(model,obj,dt,grads)

    method = :trapezoid
    method = :hermite_simpson_separated

    # Constants
    nSeg = Int(floor(obj.tf/dt)); # Number of segments
    dt = obj.tf/nSeg
    if method == :trapezoid
        N = nSeg + 1
    elseif method == :hermite_simpson_separated
        N = 2*nSeg + 1
    end
    n,m = model.n, model.m
    pack = (n,m,N)

    usrfun, eval_f, eval_g = gen_usrfun(model, obj, dt, pack, method, grads=grads)

    # Set up the problem
    Z0 = get_initial_state(obj,N)
    lb,ub = get_bounds(obj,N)

    # Set options
    options = Dict{String, Any}()
    options["Derivative option"] = 0
    options["Verify level"] = 1

    # Solve the problem
    @time z_opt, fopt, info = snopt(usrfun, Z0, lb, ub, options)
    @time snopt(usrfun, Z0, lb, ub, options)
    # xopt, fopt, info = Z0, Inf, "Nothing"
    x_opt,u_opt = unpackZ(z_opt,pack)

    println(info)
    return x_opt, u_opt, fopt
end

model = TrajectoryOptimization.Dynamics.pendulum[1]
model! = TrajectoryOptimization.Dynamics.pendulum![1]
obj = TrajectoryOptimization.Dynamics.pendulum[2]
obj = TrajectoryOptimization.ConstrainedObjective(obj, u_min=-2, u_max=2)
dt = 0.1

using Juno
# @profiler dircol(model, obj, dt, :other)
x_opt, u_opt, fopt = dircol(model, obj, dt, :none)
plot(x_opt')

weights = get_weights(:hermite_simpson_separated,N)
N = size(x_opt,2)
pack = (model.n,model.m,N)
X,U = x_opt,u_opt
Z = packZ(X,U)
f_aug! = f_augmented!(model!.f, model.n, model.m)
zdot = zeros(model.n)
F(z) = ForwardDiff.jacobian(f_aug!,zdot,z)
usrfun, eval_f, eval_g = gen_usrfun(model,obj,dt, pack, :hermite_simpson_separated, grads=:auto)

# Constraints
jacob_g = constraint_jacobian(X,U,dt,:hermite_simpson_separated,F)
jacob_g_auto = ForwardDiff.jacobian(eval_g,Z)
errors = jacob_g_auto == jacob_g

# Objective
grad_f = cost_gradient(X,U,weights,obj)
grad_f_auto = ForwardDiff.gradient(eval_f,Z)
errors = grad_f - grad_f_auto
errors[end-10:end]
grad_f ≈ grad_f_auto

#
# Z0 = get_initial_state(obj,N)
# options = Dict{String, Any}()
# lb,ub = get_bounds(obj,N)
# options["Derivative option"] = 0
# options["Verify level"] = 1
#
# @time xopt, fopt, info = snopt(gradfree, Z0, lb, ub, options)
# x_soln,u_soln = unpackZ(xopt)
# plot(x_soln')
# plot(u_soln')
# eval_g(xopt)
#
# # Tests
# X,U = unpackZ(Z0)
# X = rand(size(X))
# U = rand(size(U))
# Z = packZ(X,U)
#
# X2,U2 = unpackZ(Z)
# @test X2 == X
# @test U2 == U
#
# # Test cost
# J1 = eval_f(Z)
# J2 = cost(X,U)
# @test J1 == J2
#
# # Test Constraints
# g = eval_g(Z)
#
# # Test Gradient of Objective
# grad_f = zeros((n+m)*N)
# eval_grad_f(Z,grad_f)
# grad_f
# grad_f_auto = ForwardDiff.gradient(eval_f,Z)
# @test grad_f == grad_f_auto
#
# # Gradient of Constraints
# jacob_g = eval_jac_g(Z)
# jacob_g_auto = ForwardDiff.jacobian(eval_g,Z)
# @test jacob_g_auto == jacob_g
#
# # Test usrfun
# f_val,g2,df,dg,fail = usrfun(Z)
# @test f_val == J1
# @test g2 == g
# @test df == grad_f
# @test dg == jacob_g
# @test ~fail


text:
module Percolation

global const libpercolation = joinpath(@__DIR__, "..", "libpercolation.so")

using DelimitedFiles: readdlm, writedlm
using LinearAlgebra
using Statistics
using Base.Threads

"""
    tcm_square_lattice

Equivalent to `tcm_square_lattice_t` in "percolation.h":

```c
typedef struct tcm_square_lattice {
    int64_t (*neighbours)[4];
    int64_t length;
    int64_t size;
    bool    periodic;
} tcm_square_lattice_t;
```

It is a very low-level data type and should not be used directly. Consider using
[`SquareLattice`](@ref) instead.
"""
struct tcm_square_lattice
    neighbours::Ptr{Int64}
    length::Int64
    size::Int64
    periodic::Bool
end

"""
    tcm_cubic_lattice

Equivalent to `tcm_cubic_lattice_t` in "percolation.h":

```c
typedef struct tcm_cubic_lattice {
    int64_t (*neighbours)[6];
    int64_t length;
    int64_t size;
    bool    periodic;
} tcm_cubic_lattice_t;
```

It is a very low-level data type and should not be used directly. Consider using
[`CubicLattice`](@ref) instead.
"""
struct tcm_cubic_lattice
    neighbours::Ptr{Int64}
    length::Int64
    size::Int64
    periodic::Bool
end

"""
    tcm_triangular_lattice

Equivalent to `tcm_triangular_lattice_t` in "percolation.h":

```c
typedef struct tcm_triangular_lattice {
    int64_t (*neighbours)[6];
    int64_t length;
    int64_t length_y;
    int64_t size;
} tcm_triangular_lattice_t;
```

It is a very low-level data type and should not be used directly. Consider using
[`TriangularLattice`](@ref) instead.
"""
struct tcm_triangular_lattice
    neighbours::Ptr{Int64}
    length::Int64
    length_y::Int64
    size::Int64
end

"""
    tcm_triangular_stacked_lattice

Equivalent to `tcm_triangular_lattice_t` in "percolation.h":

```c
typedef struct tcm_triangular_stacked_lattice {
    int64_t (*neighbours)[8];
    int64_t length;
    int64_t length_y;
    int64_t size;
} tcm_triangular_stacked_lattice_t;
```

It is a very low-level data type and should not be used directly. Consider using
[`TriangularStackedLattice`](@ref) instead.
"""
struct tcm_triangular_stacked_lattice
    neighbours::Ptr{Int64}
    length::Int64
    length_y::Int64
    size::Int64
end

mutable struct tcm_random_number_generator
    raw::Ptr{Cvoid}

    function tcm_random_number_generator(n)
        p = ccall((:tcm_random_number_generator_init, libpercolation),
                  Ptr{tcm_random_number_generator}, (Cuint,), n)
        if p == C_NULL throw(OutOfMemoryError()) end
        generator = new(p)
        return finalizer(
            x -> ccall((:tcm_random_number_generator_deinit, libpercolation),
                       Cvoid, (Ptr{tcm_random_number_generator},), x.raw),
            generator)
    end
end

struct tcm_perc_results
    size::UInt32
    number_sites::Ptr{UInt32}
    number_clusters::Ptr{UInt32}
    max_cluster_size::Ptr{UInt32}
    has_wrapped_one::Ptr{UInt8}
    has_wrapped_two::Ptr{UInt8}
    chirality::Ptr{Float64}
    helicity::Ptr{Float64}
    magnetisation::Ptr{Float64}
end

struct tcm_percolation_results
    number_clusters::Ptr{Float64}
    max_cluster_size::Ptr{Float64}
    has_wrapped_one::Ptr{Float64}
    has_wrapped_two::Ptr{Float64}
    chirality::Ptr{Float64}
    helicity::Ptr{Float64}
    magnetisation::Ptr{Float64}
end

struct tcm_percolation_stats
    max_number_sites::Ptr{Float64}
    mean_number_sites::Ptr{Float64}
    max_number_children::Ptr{Float64}
    mean_number_children::Ptr{Float64}
end

struct tcm_convolution_state
    number_sites::Int64
    number_functions::Int64
    functions::Ptr{Ptr{Float64}}
    outputs::Ptr{Ptr{Float64}}
end

abstract type Lattice{N} end

"""
    SquareLattice <: Lattice{2}

A simple square lattice in two dimensions
----------------------------------------------------------------------

    SquareLattice(n, periodic=true)

Construct a new square lattice of side length `n` with (if `periodic == true`)
or without (if `periodic == false`) periodic boundary conditions.
"""
struct SquareLattice <: Lattice{2}
    neighbours::Array{Int64, 2}
    length::Int
    is_periodic::Bool

    function SquareLattice(n::Int, periodic::Bool=true)
        length_min = periodic ? 3 : 2
        if n < length_min
            throw(DomainError(n, "side length must be at least $length_min"))
        end
        lattice = new(Array{Int64, 2}(undef, 4, n^2), n, periodic)
        status = ccall((:tcm_compute_neighbours_square, libpercolation),
              Cint, (tcm_square_lattice,),
              tcm_square_lattice(pointer(lattice.neighbours), n, n^2, periodic))
        if status != 0 throw(SystemError("tcm_compute_neighbours_square", status)) end
        return lattice
    end
end

struct CubicLattice <: Lattice{3}
    neighbours::Array{Int64, 2}
    length::Int
    is_periodic::Bool

    function CubicLattice(n::Int, periodic::Bool=true)
        length_min = periodic ? 3 : 2
        if n < length_min
            throw(DomainError(n, "side length must be at least $length_min"))
        end
        lattice = new(Array{Int64, 2}(undef, 6, n^3), n, periodic)
        status = ccall((:tcm_compute_neighbours_cubic, libpercolation),
              Cint, (tcm_cubic_lattice,),
              tcm_cubic_lattice(pointer(lattice.neighbours), n, n^3, periodic))
        if status != 0 throw(SystemError("tcm_compute_neighbours_cubic", status)) end
        return lattice
    end
end

struct TriangularLattice <: Lattice{2}
    neighbours::Array{Int64, 2}
    length_x::Int
    length_y::Int

    function TriangularLattice(n::Int)
        if n < 2 throw(DomainError(n, "side length must be at least 2")) end
        length_x = n
        length_y = Int(round(n * 2 / sqrt(3.0)))
        lattice = new(Array{Int64, 2}(undef, 6, length_x * length_y), length_x, length_y)
        status = ccall((:tcm_compute_neighbours_triangular, libpercolation),
              Cint, (tcm_triangular_lattice,),
              tcm_triangular_lattice(pointer(lattice.neighbours),
                                     length_x, length_y, length_x * length_y))
        if status != 0 throw(SystemError("tcm_compute_neighbours_triangular", status)) end
        return lattice
    end
end

struct TriangularStackedLattice <: Lattice{3}
    neighbours::Array{Int64, 2}
    length::Int
    length_y::Int

    function TriangularStackedLattice(n::Int)
        if n < 2 throw(DomainError(n, "side length must be at least 2")) end
        length = n
        length_y = Int(round(n * 2 / sqrt(3.0)))
        lattice = new(Array{Int64, 2}(undef, 8, length * length * length_y),
                      length, length_y)
        status = ccall((:tcm_compute_neighbours_triangular_stacked, libpercolation),
            Cint, (tcm_triangular_stacked_lattice,),
            tcm_triangular_stacked_lattice(
                pointer(lattice.neighbours), length, length_y, length * length * length_y))
        if status != 0 throw(SystemError("tcm_compute_neighbours_triangular", status)) end
        return lattice
    end
end

number_sites(x::SquareLattice)::Int = x.length^2
number_sites(x::CubicLattice)::Int = x.length^3
number_sites(x::TriangularLattice)::Int = x.length_x * x.length_y
number_sites(x::TriangularStackedLattice)::Int = x.length * x.length * x.length_y

function tcm_lattice(x::SquareLattice)::tcm_square_lattice
    return tcm_square_lattice(
        pointer(x.neighbours), x.length, number_sites(x), x.is_periodic)
end

function tcm_lattice(x::CubicLattice)::tcm_cubic_lattice
    return tcm_cubic_lattice(
        pointer(x.neighbours), x.length, number_sites(x), x.is_periodic)
end

function tcm_lattice(x::TriangularLattice)::tcm_triangular_lattice
    return tcm_triangular_lattice(
        pointer(x.neighbours), x.length_x, x.length_y, number_sites(x))
end

function tcm_lattice(x::TriangularStackedLattice)::tcm_triangular_stacked_lattice
    return tcm_triangular_stacked_lattice(
        pointer(x.neighbours), x.length, x.length_y, number_sites(x))
end

mutable struct AverageAccumulator
    mu
    count::Int

    function AverageAccumulator(mu)
        fill!(mu, 0)
        return new(mu, 0)
    end
end

function call!(acc::AverageAccumulator, x)
    @assert length(acc.mu) == length(x)
    acc.count += 1
    BLAS.axpy!(-1.0, acc.mu, x)
    BLAS.axpy!(1.0 / acc.count, x, acc.mu)
end

function _percolate_impl(lattice::tcm_square_lattice,
                         out::tcm_perc_results,
                         rng::tcm_random_number_generator)::Nothing
    status = ccall((:tcm_percolate_square, libpercolation),
                   Cint, (Ref{tcm_square_lattice},
                          Ref{tcm_perc_results},
                          Ptr{tcm_percolation_stats},
                          Ptr{tcm_random_number_generator}),
                   lattice, out, C_NULL, rng.raw)
    if status != 0 throw(SystemError("tcm_percolate_square", status)) end
end

function percolate_test(xs, lattice::Lattice{N};
                        batch_size=1, out::Union{AbstractArray{Float64, 2}, Nothing}=nothing) where {N}
    if out == nothing out = Array{Float64, 2}(undef, length(xs), 4) end
    fill!(out, 0.0)

    number_clusters = Array{UInt32, 1}(undef, length(xs))
    max_cluster_size = Array{UInt32, 1}(undef, length(xs))
    has_wrapped_one = Array{UInt8, 1}(undef, length(xs))
    has_wrapped_two = Array{UInt8, 1}(undef, length(xs))

    results = tcm_perc_results(
        UInt32(length(xs)),
        pointer(xs),
        pointer(number_clusters),
        pointer(max_cluster_size),
        pointer(has_wrapped_one),
        pointer(has_wrapped_two),
        C_NULL,
        C_NULL,
        C_NULL)
    rng = tcm_random_number_generator(123)

    for i in 1:batch_size
        _percolate_impl(tcm_lattice(lattice), results, rng)
        view(out, :, 1) .+= number_clusters
        view(out, :, 2) .+= max_cluster_size
        view(out, :, 3) .+= has_wrapped_one
        view(out, :, 4) .+= has_wrapped_two
    end
    out ./= batch_size
    return out
end

function _percolate(n_min::Csize_t, n_max::Csize_t,
                    lattice::tcm_square_lattice,
                    out::tcm_percolation_results,
                    stats::Union{tcm_percolation_stats, Nothing})::Nothing
    status = ccall((:tcm_percolate_square, libpercolation),
        Cint, (Csize_t, Csize_t, tcm_square_lattice, Ref{tcm_percolation_results}, Ptr{tcm_percolation_stats}),
        n_min, n_max, lattice, out, stats != nothing ? Ref(stats) : C_NULL)
    if status != 0 throw(SystemError("tcm_percolate_square", status)) end
end

function _percolate(lattice::tcm_square_lattice, out::tcm_percolation_results)::Nothing
    status = ccall((:tcm_percolate_square, libpercolation),
        Cint, (tcm_square_lattice, Ref{tcm_percolation_results}, Ptr{tcm_percolation_stats}),
        lattice, out, C_NULL)
    if status != 0 throw(SystemError("tcm_percolate_square", status)) end
end

function _percolate(lattice::tcm_cubic_lattice, out::tcm_percolation_results)::Nothing
    status = ccall((:tcm_percolate_cubic, libpercolation),
        Cint, (tcm_cubic_lattice, Ref{tcm_percolation_results}), lattice, out)
    if status != 0 throw(SystemError("tcm_percolate_cubic", status)) end
end

function _percolate(lattice::tcm_triangular_lattice, out::tcm_percolation_results)::Nothing
    status = ccall((:tcm_percolate_triangular, libpercolation),
        Cint, (tcm_triangular_lattice, Ref{tcm_percolation_results}), lattice, out)
    if status != 0 throw(SystemError("tcm_percolate_triangular", status)) end
end

function _percolate(lattice::tcm_triangular_stacked_lattice, out::tcm_percolation_results)::Nothing
    status = ccall((:tcm_percolate_triangular_stacked, libpercolation),
        Cint, (tcm_triangular_stacked_lattice, Ref{tcm_percolation_results}), lattice, out)
    if status != 0 throw(SystemError("tcm_percolate_triangular", status)) end
end

function percolate(n_min, n_max, lattice::Lattice{N}, batch_size=1, profile=false;
                   out::Union{AbstractArray{Float64, 2}, Nothing}=nothing) where {N}
    if batch_size < 1 throw(DomainError(batch_size, "batch size must be at least 1")) end
    if N != 2
        throw(TypeError(:percolate,
            "only 2-dimensional lattices are supported.", Lattice))
    end
    n = number_sites(lattice)
    number_columns = 4
    if out == nothing out = Array{Float64, 2}(undef, n + 1, number_columns) end
    @assert size(out) == (n + 1, number_columns) # NOTE: This +1 is very important!!!
    acc  = AverageAccumulator(out)
    local_result = Array{Float64, 2}(undef, n + 1, number_columns)
    results = tcm_percolation_results(
        pointer(view(local_result, :, 1)),
        pointer(view(local_result, :, 2)),
        pointer(view(local_result, :, 3)),
        C_NULL,
        C_NULL,
        C_NULL,
        pointer(view(local_result, :, 4)),
    )

    if profile
        stats = Array{Float64, 2}(undef, n + 1, 4)
        stats_acc = AverageAccumulator(stats)
        local_stats = Array{Float64, 2}(undef, n + 1, 4)
        tcm_stats = tcm_percolation_stats(
            pointer(view(local_stats, :, 1)),
            pointer(view(local_stats, :, 2)),
            pointer(view(local_stats, :, 3)),
            pointer(view(local_stats, :, 4))
        )
    else
        tcm_stats = nothing
    end
    for i in 1:batch_size
        _percolate(Csize_t(n_min), Csize_t(n_max), tcm_lattice(lattice), results, tcm_stats)
        call!(acc, local_result)
        if profile call!(stats_acc, local_stats) end
    end
    fill!(view(out, 1:max(n_min-1, 1), :), 0.0)
    fill!(view(out, min(n_max+1, n + 1):(n + 1), :), 0.0)
    return profile ? (out, stats) : out
end

function percolate_parallel(lattice::Lattice{N}, batch_size=1) where {N}
    if batch_size < 1 throw(DomainError(batch_size, "batch size must be at least 1")) end
    if N != 2 && N != 3
        throw(TypeError(:percolate,
            "only 2- and 3-dimensional lattices are supported.", Lattice))
    end
    n = number_sites(lattice)
    number_workers = nthreads()
    number_columns = N == 2 ? 3 : 4
    @assert mod(batch_size, number_workers) == 0
    local_batch_size = div(batch_size, number_workers)
    results = Array{Float64, 3}(undef, n + 1, number_columns, number_workers)
    @threads for i in 1:number_workers
        percolate(lattice, local_batch_size; out=view(results, :, :, i))
    end
    mean = Array{Float64, 2}(undef, n + 1, number_columns)
    return Statistics.mean!(mean, results)
end

function convolution(measurements::AbstractArray{Float64, 2};
                     p_min::Float64=0.0, p_max::Float64=1.0)
    number_sites = size(measurements, 1) - 1
    N = size(measurements, 2)
    # Some sanity checks
    @assert number_sites > 0 "arrays must contain at least 2 elements"
    @assert 0.0 <= p_min && p_min <= p_max && p_max <= 1.0
    n_min = Int(floor(p_min * number_sites))
    n_max = Int(floor(p_max * number_sites))
    @assert n_min <= n_max
    result_size = n_max - n_min + Int32(1)
    # Constructing the convolution state
    results = Array{Float64, 2}(undef, result_size, N + 1)
    funcs::Array{Ptr{Float64}, 1} = [pointer(view(measurements, :, i)) for i in 1:N]
    outs::Array{Ptr{Float64}, 1} =
        [pointer(view(results, :, i)) - n_min * sizeof(Float64) for i in 1:N+1]
    state = tcm_convolution_state(
        Int64(number_sites), Int64(N), pointer(funcs), pointer(outs))
    status = ccall((:tcm_convolution, libpercolation),
        Cint, (Int64, Int64, Ref{tcm_convolution_state},),
        n_min, n_max, state)
    if status != 0 throw(SystemError("tcm_convolution", status)) end
    return results
end

# function convolution(measurements::NTuple{N, Array{Float64, 1}};
#                      p_min::Float64=0.0, p_max::Float64=1.0) where {N}
#     # Trivial case
#     if N == 0 return () end
#     number_sites = length(measurements[1]) - 1
#     # Some sanity checks
#     @assert number_sites > 0 "arrays must contain at least 2 elements"
#     for m in measurements
#         @assert length(m) == number_sites + 1 "all arrays must be of the same length"
#     end
#     @assert 0.0 <= p_min && p_min <= p_max && p_max <= 1.0
#     n_min = Int32(floor(p_min * number_sites))
#     n_max = Int32(floor(p_max * number_sites))
#     @assert n_min <= n_max
#     result_size = n_max - n_min + Int32(1)
#     # Constructing the convolution state
#     results::NTuple{N + 1, Array{Float64, 1}} = (
#         Array{Float64, 1}(undef, result_size),
#         [Array{Float64, 1}(undef, result_size) for m in measurements]...)
#     fns::Array{Ptr{Float64}, 1} = [pointer(m) for m in measurements]
#     outs::Array{Ptr{Float64}, 1} = [pointer(r) - n_min * sizeof(Float64) for r in results]
#     state = tcm_convolution_state(
#         Int32(number_sites), Int32(N), pointer(fns), pointer(outs))
#     status = ccall((:tcm_convolution, libpercolation),
#         Cint, (Int32, Int32, Ref{tcm_convolution_state},),
#         n_min, n_max, state)
#     if status != 0 throw(SystemError("tcm_convolution", status)) end
#     return results
# end

end # module

max_stars_repo_name: twesterhout/percolation

max_issues_repo_name: twesterhout/percolation

max_forks_repo_name: twesterhout/percolation


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct ErrorResponse <: SwaggerModel
    code::Any # spec type: Union{ Nothing, String } # spec name: code
    message::Any # spec type: Union{ Nothing, String } # spec name: message
    target::Any # spec type: Union{ Nothing, String } # spec name: target
    details::Any # spec type: Union{ Nothing, Vector{ErrorResponse} } # spec name: details
    additionalInfo::Any # spec type: Union{ Nothing, Vector{ErrorAdditionalInfo} } # spec name: additionalInfo

    function ErrorResponse(;code=nothing, message=nothing, target=nothing, details=nothing, additionalInfo=nothing)
        o = new()
        validate_property(ErrorResponse, Symbol("code"), code)
        setfield!(o, Symbol("code"), code)
        validate_property(ErrorResponse, Symbol("message"), message)
        setfield!(o, Symbol("message"), message)
        validate_property(ErrorResponse, Symbol("target"), target)
        setfield!(o, Symbol("target"), target)
        validate_property(ErrorResponse, Symbol("details"), details)
        setfield!(o, Symbol("details"), details)
        validate_property(ErrorResponse, Symbol("additionalInfo"), additionalInfo)
        setfield!(o, Symbol("additionalInfo"), additionalInfo)
        o
    end
end # type ErrorResponse

const _property_map_ErrorResponse = Dict{Symbol,Symbol}(Symbol("code")=>Symbol("code"), Symbol("message")=>Symbol("message"), Symbol("target")=>Symbol("target"), Symbol("details")=>Symbol("details"), Symbol("additionalInfo")=>Symbol("additionalInfo"))
const _property_types_ErrorResponse = Dict{Symbol,String}(Symbol("code")=>"String", Symbol("message")=>"String", Symbol("target")=>"String", Symbol("details")=>"Vector{ErrorResponse}", Symbol("additionalInfo")=>"Vector{ErrorAdditionalInfo}")
Base.propertynames(::Type{ ErrorResponse }) = collect(keys(_property_map_ErrorResponse))
Swagger.property_type(::Type{ ErrorResponse }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_ErrorResponse[name]))}
Swagger.field_name(::Type{ ErrorResponse }, property_name::Symbol) =  _property_map_ErrorResponse[property_name]

function check_required(o::ErrorResponse)
    true
end

function validate_property(::Type{ ErrorResponse }, name::Symbol, val)
end


text:
__precompile__(true)
module SuperHF

include("HF.jl")
include("HF_Hubbard.jl")

end ## End of module

max_stars_repo_name: oliviersimard/Hartree-Fock

max_issues_repo_name: oliviersimard/Hartree-Fock

max_forks_repo_name: oliviersimard/Hartree-Fock


text:
export
   maximize_quantum_search

"""
    maximize_quantum_search(qws_cont [, maxtime, tstep])

Determines optimal runtime for continuous quantum walk search. The time is
searched in [0, maxtime] interval, with penalty `penalty(qws_cont)`.
It is recommended for penalty to be nonzero, otherwise time close to 0 is usually
returned. Typically `penalty` equal to `log(nv(graph(qws_cont)))` is enough, but
optimal value may depend on the model or graph chosen.

The optimal time is chosen according to expected runtime, which equals to
runtime over probability. This comes from interpreting the evolution as the
Bernoulli process.

`tstep` is used for primary grid search to search for determine intervale which
is supsected to have small expected runtime. To large value may miss the optimal value,
while to small may greatly increase runtime of the algorithm.

`maxtime` defaults to graph order n, `tstep` defaults to `sqrt(n)/5`. Note that in general the probability is not maximal.

```jldoctest; setup = :(using QuantumWalk,LightGraphs)
julia> qws = QWSearch(CTQW(complete_graph(100)), [1], 1., 0.01);

julia> result = maximize_quantum_search(qws)
QSearchState{Array{Complex{Float64},1},Float64,Float64}(Complex{Float64}[0.6211421326794925 + 0.6956647632995185im, 0.027973603660103912 - 0.023085974674888862im, 0.027973603660103912 - 0.023085974674888862im, 0.027973603660103912 - 0.023085974674888862im, 0.027973603660103912 - 0.023085974674888862im, 0.027973603660103912 - 0.023085974674888862im, 0.027973603660103912 - 0.023085974674888862im, 0.027973603660103912 - 0.023085974674888862im, 0.027973603660103912 - 0.023085974674888862im, 0.027973603660103912 - 0.023085974674888862im  …  0.027973603660103912 - 0.023085974674888862im, 0.027973603660103912 - 0.023085974674888862im, 0.027973603660103912 - 0.023085974674888862im, 0.027973603660103912 - 0.023085974674888862im, 0.027973603660103912 - 0.023085974674888862im, 0.027973603660103912 - 0.023085974674888862im, 0.027973603660103895 - 0.023085974674888883im, 0.027973603660103895 - 0.023085974674888883im, 0.027973603660103895 - 0.023085974674888883im, 0.027973603660103895 - 0.023085974674888883im], [0.8697670118862033], 11.996369403338626, 1.0)

julia> probability(result)
1-element Array{Float64,1}:
 0.8697670118862033
```
"""
function maximize_quantum_search(qws::QWSearch{<:QWModelCont},
                                 maxtime::T = Float64(nv(graph(qws))),
                                 tstep::T = Float64(0.2*sqrt(nv(graph(qws))))) where T<:Real
   @assert maxtime >= 0. "Parameter 'maxtime' needs to be nonnegative"
   @assert tstep > 0. "Parameter 'tstep' needs to be positive"

   if penalty(qws) == 0
      @warn "It is recommended for the penalty to be nonzero. Otherwise, the time close to zero may be returned. "*
      "Typically penalty should be approximately log(n), but this might by case-dependant."
   end

   state = initial_state(qws)
   function efficiency_opt(runtime::Number)
      expected_runtime(runtime+penalty(qws), sum(measure(qws, evolve(qws, state, runtime), marked(qws))))
   end

   t = zero(T)
   data_t = T[t]
   data_y = [efficiency_opt(t)]
   max_efficiency = data_y[end]
   for t=tstep:tstep:maxtime
      push!(data_t, t)
      push!(data_y, efficiency_opt(t))

      max_efficiency = min(max_efficiency, data_y[end])
      if max_efficiency <= t
         break
      end
   end

   if t > maxtime-tstep
      push!(data_t, maxtime)
      push!(data_y, efficiency_opt(maxtime))
   end

   minindex = findmin(data_y)[2]

   mint = max(zero(T), data_t[max(1, minindex-1)])
   maxt = min(maxtime, data_t[min(length(data_t), minindex+1)])

   optresult = optimize(efficiency_opt, mint, maxt)
   result = execute_single(qws, Optim.minimizer(optresult))
   QSearchState(result.state, result.probability, result.runtime, penalty(qws))
end

"""
    maximize_quantum_search(qws_discr [, runtime, mode])

Determines optimal runtime for discrete quantum walk search. The time is
searched in [0, runtime] interval, with penalty `penalty(qws_discr)`.
It is recommended for penalty to be nonzero, otherwise time close 0 is returned.
Typically small `penalty` approximately equal to log(n) is enough, but
optimal value may depend on the model or graph chosen.

The optimal time depends on chosen `mode`:
* `:firstmaxprob` stops when probability start to decrease,
* `:firstmaxeff` stops when expected runtime start to increase,
* `:maxtimeeff` chooses exhaustively the time from [0, runtime] with smallest expected time,
* `:maxtimeprob` chooses exhaustively the time from [0, runtime] with maximal success probability,
* `:maxeff` (default) finds optimal time with smallest expected time, usually faster than `:maxtimefff`.

Note last three modes always returns optimal time within the interval.

`maxtime` defaults to graph order n, `mode` defaults to `:maxeff`.

```jldoctest; setup = :(using QuantumWalk,LightGraphs)
julia> qws = QWSearch(Szegedy(complete_graph(200)), [1], 1);

julia> result = maximize_quantum_search(qws);

julia> runtime(result)
6

julia> probability(result)
1-element Array{Float64,1}:
 0.5000160413993847

julia> result = maximize_quantum_search(qws, 100, :maxtimeprob);

julia> runtime(result)
39

julia> probability(result)
1-element Array{Float64,1}:
 0.5509378780415133
```
"""
function maximize_quantum_search(qws::QWSearch{<:QWModelDiscr},
                                 runtime::Int = nv(graph(qws)),
                                 mode::Symbol = :maxeff)
   @assert runtime>=0 "Parameter 'runtime' needs to be nonnegative"
   @assert mode ∈ [:firstmaxprob, :firstmaxeff, :maxtimeeff, :maxeff, :maxtimeprob] "Specified stop condition is not implemented"
   if penalty(qws) == 0
      @warn "It is recommended for the penalty to be nonzero. Otherwise, the time close to zero is returned. "*
      "Typically penalty should be approximately log(n), but this might by case-dependant."
   end

   best_result = QSearchState(qws, initial_state(qws), 0)
   state = QSearchState(qws, initial_state(qws), 0)
   for t=1:runtime
      state = QSearchState(qws, evolve(qws, state), t)
      stopsearchflag = stopsearch(qws, best_result, state, mode)
      best_result = best(qws, best_result, state, mode)

      if stopsearchflag
         break
      end
   end

   best_result
end

"""
    stopsearch(qws, previous_state, state, mode)

For given combination of argument decides whetver maximizing search function
should be stopped:
The optimal time depende on chosen `mode`:
* `:firstmaxprob` stops when probability start to decrease,
* `:firstmaxeff` stops whene expected runtime start to increase,
* `:maxtimeeff` always return `false` (should be decided by external function),
* `:maxtimeprob`  always return `false` (should be decided by external function),
* `:maxeff` checks whetver expected_runtime of older state is smaller than current time,

`false` means maximizing should continue unless other constraints, `true` otherwise.
"""
function stopsearch(qws::QWSearch,
                    prev_state::QSearchState,
                    state::QSearchState,
                    mode::Symbol)
   if mode == :maxeff
      return expected_runtime(qws, prev_state) < state.runtime+1 #check whetver needs to analyze next step
   elseif mode == :firstmaxprob
      return sum(prev_state.probability) > sum(state.probability)
   elseif mode == :firstmaxeff
      return expected_runtime(qws, prev_state) < expected_runtime(qws, state)
   else # include :maxtime case, should be considered by outside loop (hack?)
      return false
   end
end

"""
    best(qws, state1, state2, mode)

Choose better state depending on mode. If success probability is maximized, then
the success probability is compared. If expected runtime is minimized,
then expected runtime is compared.
"""
function best(qws::QWSearch,
              state1::QSearchState,
              state2::QSearchState,
              mode::Symbol)
   if mode ∈ [:firstmaxprob,:maxtimeprob]
      sum(state1.probability) > sum(state2.probability) ? state1 : state2
   else
      expected_runtime(qws, state1) < expected_runtime(qws, state2) ? state1 : state2
   end
end


text:
using Test
using RDates
using Dates
using ParserCombinator

@testset "rdate calendar adjustments" begin
    cal_mgr = RDates.SimpleCalendarManager(Dict("WEEKEND" => RDates.WeekendCalendar()))
    with_cal_mgr(cal_mgr) do
        @test RDates.CalendarAdj(["WEEKEND"], rd"1d", RDates.HolidayRoundingNBD()) + Date(2019,4,16) == Date(2019,4,17)
        @test RDates.CalendarAdj(["WEEKEND"], rd"1d", RDates.HolidayRoundingNBD()) + Date(2019,9,27) == Date(2019,9,30)
        @test RDates.CalendarAdj(["WEEKEND"], rd"1d", RDates.HolidayRoundingNBD()) + Date(2019,11,29) == Date(2019,12,2)
        @test RDates.CalendarAdj(["WEEKEND"], rd"1d", RDates.HolidayRoundingNBDSM()) + Date(2019,11,29) == Date(2019,11,29)
    end
end

@testset "rdate business days" begin
    cal_mgr = RDates.SimpleCalendarManager(Dict("WEEKEND" => RDates.WeekendCalendar()))
    with_cal_mgr(cal_mgr) do
        @test rd"0b@WEEKEND" + Date(2019,9,28) == Date(2019,9,30)
        @test rd"-0b@WEEKEND" + Date(2019,9,28) == Date(2019,9,27)
        @test rd"1b@WEEKEND" + Date(2019,9,28) == Date(2019,9,30)
        @test rd"2b@WEEKEND" + Date(2019,9,28) == Date(2019,10,1)
    end
end

@testset "rdate add ordering" begin
    @test rd"1d" + Date(2019,4,16) == Date(2019,4,17)
    @test Date(2019,4,16) + rd"1d" == Date(2019,4,17)
    @test rd"1d+2d" == rd"1d" + rd"2d"
    @test rd"3*1d" == RDates.multiply_roll(rd"1d",3)
end

@testset "different month and year conventions" begin
    @test RDates.Year(1) + Date(2016,2,29) == Date(2017,2,28)
    @test Dates.Year(1) + Date(2016,2,29) == Date(2017,2,28)
    @test rd"1y" + Date(2016,2,29) == Date(2017,2,28)
    @test RDates.Year(1, RDates.InvalidDayLDOM(), RDates.MonthIncrementPDOM()) + Date(2016,2,29) == Date(2017,2,28)
    @test rd"1y[LDOM;PDOM]" + Date(2016,2,29) == Date(2017,2,28)
    @test RDates.Year(1, RDates.InvalidDayFDONM(), RDates.MonthIncrementPDOM()) + Date(2016,2,29) == Date(2017,3,1)
    @test rd"1y[FDONM;PDOM]" + Date(2016,2,29) == Date(2017,3,1)

    @test RDates.Month(1) + Date(2019,1,31) == Date(2019,2,28)
    @test Dates.Month(1) + Date(2019,1,31) == Date(2019,2,28)
    @test rd"1m" + Date(2019,1,31) == Date(2019,2,28)
    @test RDates.Month(1, RDates.InvalidDayNDONM(), RDates.MonthIncrementPDOM()) + Date(2019,1,31) == Date(2019,3,3)
    @test rd"1m[NDONM;PDOM]" + Date(2019,1,31) == Date(2019,3,3)
    @test RDates.Month(1, RDates.InvalidDayFDONM(), RDates.MonthIncrementPDOM()) + Date(2019,1,31) == Date(2019,3,1)
    @test rd"1m[FDONM;PDOM]" + Date(2019,1,31) == Date(2019,3,1)
    @test RDates.Month(1, RDates.InvalidDayNDONM(), RDates.MonthIncrementPDOMEOM()) + Date(2019,1,31) == Date(2019,2,28)
    @test rd"1m[NDONM;PDOMEOM]" + Date(2019,1,31) == Date(2019,2,28)
    @test RDates.Month(1, RDates.InvalidDayNDONM(), RDates.MonthIncrementPDOMEOM()) + Date(2019,1,30) == Date(2019,3,2)
    @test rd"1m[NDONM;PDOMEOM]" + Date(2019,1,30) == Date(2019,3,2)

    # PDOMEOM can also support calendars to work off the last business day of the month
    cal_mgr = RDates.SimpleCalendarManager(Dict("WEEKEND" => RDates.WeekendCalendar()))
    with_cal_mgr(cal_mgr) do
        @test rd"3m[LDOM;PDOMEOM@WEEKEND]" + Date(2019,8,30) == Date(2019,11,29)
    end
end

@testset "rdate whitespace" begin
    @test rd"  1d" == rd"1d"
    @test rd"1d  + 3d" == rd"1d+3d"
    @test rd"--1d" == rd"1d"
end

@testset "relative add days" begin
    @test rd"1d" + Date(2019,4,16) == Date(2019,4,17)
    @test rd"1d" + Date(2019,4,30) == Date(2019,5,1)
    @test rd"0d" + Date(2015,3,23) == Date(2015,3,23)
    @test rd"7d" + Date(2017,10,25) == Date(2017,11,1)
    @test rd"-1d" + Date(2014,1,1) == Date(2013,12,31)
end

@testset "relative add weeks" begin
    @test rd"1w" + Date(2019,4,16) == Date(2019,4,23)
    @test rd"1w" + Date(2019,4,30) == Date(2019,5,7)
    @test rd"0w" + Date(2015,3,23) == Date(2015,3,23)
    @test rd"7w" + Date(2017,10,25) == Date(2017,12,13)
    @test rd"-1w" + Date(2014,1,1) == Date(2013,12,25)
end

@testset "relative add months" begin
    @test rd"1m" + Date(2019,4,16) == Date(2019,5,16)
    @test rd"1m" + Date(2019,4,30) == Date(2019,5,30)
    @test rd"0m" + Date(2015,3,23) == Date(2015,3,23)
    @test rd"12m" + Date(2017,10,25) == Date(2018,10,25)
    @test rd"-1m" + Date(2014,1,1) == Date(2013,12,1)
end

@testset "relative add years" begin
    @test rd"1y" + Date(2019,4,16) == Date(2020,4,16)
    @test rd"1y" + Date(2019,4,30) == Date(2020,4,30)
    @test rd"0m" + Date(2015,3,23) == Date(2015,3,23)
    @test rd"12y" + Date(2017,10,25) == Date(2029,10,25)
    @test rd"-1y" + Date(2014,1,1) == Date(2013,1,1)
end

@testset "relative add daymonths" begin
    @test rd"12MAR" + Date(2018,3,3) == Date(2018,3,12)
    @test rd"1JAN" + Date(2019,12,31) == Date(2019,1,1)
    @test rd"1JAN" + Date(2020,1,1) == Date(2020,1,1)
end

@testset "relative add easter" begin
    @test rd"0E" + Date(2018,3,3) == Date(2018,4,1)
    @test rd"0E" + Date(2018,12,3) == Date(2018,4,1)
    @test rd"1E" + Date(2018,3,3) == Date(2019,4,21)
    @test rd"-1E" + Date(2018,3,3) == Date(2017,4,16)
end

@testset "relative add weekdays" begin
    @test rd"1MON" + Date(2017,10,25) == Date(2017,10,30)
    @test rd"10SAT" + Date(2017,10,25) == Date(2017,12,30)
    @test rd"-1WED" + Date(2017,10,25) == Date(2017,10,18)
    @test rd"-10FRI" + Date(2017,10,25) == Date(2017,8,18)
    @test rd"-1TUE" + Date(2017,10,25) == Date(2017,10,24)
end

@testset "dates as rdates" begin
    @test rd"1JAN2019" + Date(2017,10,25) == Date(2019,1,1)
    @test RDates.Date(Date(2019,1,1)) + Date(2017,10,25) == Date(2019,1,1)
end

@testset "negations for 'non-supported' rdates" begin
    @test rd"1d" - rd"1st MON" == rd"1d" + rd"1st MON"
end

@testset "relative add nth weekdays" begin
    @test rd"1st MON" + Date(2017,10,25) == Date(2017,10,2)
    @test rd"2nd FRI" + Date(2017,10,25) == Date(2017,10,13)
    @test rd"4th SAT" + Date(2017,11,25) == Date(2017,11,25)
    @test rd"5th SUN" + Date(2017,12,25) == Date(2017,12,31)
end

@testset "relative add nth last weekdays" begin
    @test rd"Last MON" + Date(2017,10,24) == Date(2017,10,30)
    @test rd"2nd Last FRI" + Date(2017,10,24) == Date(2017,10,20)
    @test rd"5th Last SUN" + Date(2017,12,24) == Date(2017,12,3)
end

@testset "bad nth weekdays" begin
    @test_throws ArgumentError rd"5th WED" + Date(2017,10,25)
    @test_throws ArgumentError rd"5th MON" + Date(2017,6,1)
end

@testset "relative begin or end of month" begin
    @test rd"FDOM" + Date(2017,10,25) == Date(2017,10,1)
    @test rd"LDOM" + Date(2017,10,25) == Date(2017,10,31)

    # FDOM and LDOM can support calendars as well.
    cal_mgr = RDates.SimpleCalendarManager(Dict("WEEKEND" => RDates.WeekendCalendar()))
    with_cal_mgr(cal_mgr) do
        @test rd"FDOM@WEEKEND" + Date(2019,9,25) == Date(2019,9,2)
        @test rd"LDOM@WEEKEND" + Date(2019,8,25) == Date(2019,8,30)
    end
end

@testset "basic addition and multiplication compounds" begin
    @test rd"1d+1d" + Date(2017,10,26) == Date(2017,10,28)
    @test rd"2*1d" + Date(2017,10,26) == Date(2017,10,28)
    @test rd"1d+1d+1d+1d" + Date(2017,10,26) == Date(2017,10,30)
    @test rd"4*1d" + Date(2017,10,26) == Date(2017,10,30)
    @test rd"2*2d" + Date(2017,10,26) == Date(2017,10,30)
    @test rd"1d-1d+1d-1d" + Date(2017,10,26) == Date(2017,10,26)
    @test rd"2*3d+1d" + Date(2019,4,16) == Date(2019,4,23)
    @test rd"2*(3d+1d)" + Date(2019,4,16) == Date(2019,4,24)
    @test rd"2d-1E" + Date(2019,4,16) == Date(2018,4,1)
    @test rd"1d" - rd"2*1d" + Date(2019,5,1) == Date(2019,4,30)
    @test RDates.multiply_no_roll(rd"1d", 3) + Date(2017,4,14) == Date(2017,4,17)
end

@testset "roll versus no roll addition" begin
    # If we apply each 1m addition individually, then invalid days in Feb kicks in.
    @test rd"1m+1m" + Date(2019,1,31) == Date(2019,3,28)
    @test rd"2*roll(1m)" + Date(2019,1,31) == Date(2019,3,28)
    # Normal multiplication will embed this in the period instead.
    @test rd"2m" + Date(2019,1,31) == Date(2019,3,31)
    @test rd"2*1m" + Date(2019,1,31) == Date(2019,3,31)
end

@testset "rdate ranges" begin
    @test collect(range(Date(2017,1,1), Date(2017,1,3), rd"1d")) == [Date(2017,1,1), Date(2017,1,2), Date(2017,1,3)]
    @test collect(range(Date(2017,1,1), Date(2017,1,3), rd"2d")) == [Date(2017,1,1), Date(2017,1,3)]
    @test collect(range(Date(2017,1,1), Date(2017,1,18), rd"1d+1w")) == [Date(2017,1,1), Date(2017,1,9), Date(2017,1,17)]
    @test collect(range(Date(2017,1,1), Date(2017,1,3), rd"1d", inc_from=false)) == [Date(2017,1,2), Date(2017,1,3)]
    @test collect(range(Date(2017,1,1), Date(2017,1,3), rd"1d", inc_to=false)) == [Date(2017,1,1), Date(2017,1,2)]
    @test collect(range(Date(2017,1,1), Date(2017,1,3), rd"1d", inc_from=false, inc_to=false)) == [Date(2017,1,2)]
    @test collect(range(Date(2017,1,1), Date(2018,1,1), rd"1DEC2017 + 3m + 3rd WED")) == [Date(2017,3,15), Date(2017,6,21), Date(2017,9,20), Date(2017,12,20)]
end

@testset "rdate parsing methods" begin
    @test rdate("1d") == rd"1d"
    @test rdate("3*2d") == rd"3*2d"
end

@testset "fail to parse bad rdates" begin
    @test_throws ParserCombinator.ParserException rdate("1*2")
    @test_throws ParserCombinator.ParserException rdate("1dw")
    @test_throws ParserCombinator.ParserException rdate("+2d")
    @test_throws ParserCombinator.ParserException rdate("2d+")
    @test_throws ParserCombinator.ParserException rdate("d")
end

@testset "string form" begin
    @test string(rd"1d") == "1d"
    @test string(rd"3d + 2w") == "17d"
    @test string(rd"2*(1m + 1y)") == "2m[LDOM;PDOM]+2y[LDOM;PDOM]"
    @test string(rd"2*roll(1m)") == "2*roll(1m[LDOM;PDOM])"
    @test string(rd"1m[LDOM;PDOMEOM@A]@B[NBD]") == "1m[LDOM;PDOMEOM@A]@B[NBD]"
end


text:
# ---
# title: 22. Generate Parentheses
# id: problem22
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: String, Backtracking
# link: <https://leetcode.com/problems/generate-parentheses/description/>
# hidden: true
# ---
#
# Given `n` pairs of parentheses, write a function to _generate all combinations
# of well-formed parentheses_.
#
#
#
# **Example 1:**
#
#
#
#     Input: n = 3
#     Output: ["((()))","(()())","(())()","()(())","()()()"]
#
#
# **Example 2:**
#
#
#
#     Input: n = 1
#     Output: ["()"]
#
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 8`
#
#
## @lc code=start
using LeetCode

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


text:
@reexport module EmptyReservoirs

# Standard library
using Statistics
using Printf
using Dates

using RecipesBase
using Distributions
using Parameters
using StochasticPrograms
using Flux
using HydroModels
using HydroModels: AbstractHydroModel, River, Plant, Area, Scenario, Forecaster, forecast

import HydroModels: modelindices

using Plots: font, text, Shape

export
    EmptyReservoirsData,
    EmptyReservoirsScenario,
    EmptyReservoirsSampler,
    EmptyReservoirsModel

include("data.jl")
include("scenarios.jl")
include("indices.jl")
include("model.jl")

end


text:
mutable struct Problem
  F::Firm

  N::Int
  S::Int
  p::Matrix{Float64}

  dtol::Float64
  maxiter::Int
end

function updatefirmτ!(Φ::Problem, τ::Vector{Float64})::Problem
  Φ.F.τ .= τ
  refreshR!(Φ.F, Φ.p)

  return Φ
end

function findpolicy!(Φ::Problem)::Problem
  local rslices::Vector{SubArray} #view of r in each state
  local vslices::Vector{SubArray} #view of v in each state
  local πslices::Vector{SubArray} #view of the policy in each state
  local mslices::Matrix{SubArray} #view of the pricing kernel in each state and each future state
  local iter::Int = 0
  local derr::Float64 = 1.0
  local vrow::Matrix{Float64} = Matrix{Float64}(undef,1,Φ.N) #captures updated values
  local πrow::Matrix{CartesianIndex} = Matrix{CartesianIndex}(undef,1,Φ.N) #captures updated policy
  local vold::Matrix{Float64} = deepcopy(Φ.F.V) #holds the previous value for comparison purposes

  #each view here maps all current states to a single future state
  rslices = Vector{SubArray}(undef, Φ.S)
  mslices = Matrix{SubArray}(undef, Φ.S, Φ.S)
  vslices = Vector{SubArray}(undef, Φ.S)
  πslices = Vector{SubArray}(undef, Φ.S)

  #make our convenience views
  for s ∈ 1:Φ.S
    rslices[s] = view(Φ.F.R, :, :,  s)
    vslices[s] = view(Φ.F.V, :, s)
    πslices[s] = view(Φ.F.Π, :, s)
    mslices[s,:] .= ((sₜ₁::Int)->view(Φ.F.M, :, :, s, sₜ₁)).(1:Φ.S)
  end

  #display(rslices[1] .+
  #  Φ.p[1, 1] .* mslices[1,1] .* vslices[1] .+
  #  Φ.p[2, 1] .* mslices[1,2] .* vslices[2] .+
  #  Φ.p[3, 1] .* mslices[1,3] .* vslices[3])

  while (derr > Φ.dtol) && (iter < Φ.maxiter)
    for s ∈ 1:Φ.S #need to use vec since findmax returns column vectors
      vrow, πrow = findmax(rslices[s] .+
        Φ.p[1, s] .* mslices[s,1] .* vslices[1] .+
        Φ.p[2, s] .* mslices[s,2] .* vslices[2] .+
        Φ.p[3, s] .* mslices[s,3] .* vslices[3], dims=1)

      #need to do this to fix the dimensionality of the returned values from findmax
      vslices[s] .= vec(vrow)
      πslices[s] .= vec(πrow)
      #isplay(vrow)
      #=println(maximum(Φ.p[1, s] .* mslices[s,1] .+
        Φ.p[2, s] .* mslices[s,2] .+
        Φ.p[3, s] .* mslices[s,3]))=#
    end

    derr = maximum((abs).(vold .- Φ.F.V))
    vold .= Φ.F.V

    iter += 1

    (iter % 100 == 0) && println("iter: $iter, derr: $derr")

  end

  (iter == Φ.maxiter) && @warn("Convergence failure iter: $iter, derr: $derr")

  return Φ
end

function Problem(F::Firm{NTuple{3,Float64}}, p::Matrix{Float64};
    dtol::Float64=DEF_D_TOL, maxiter::Int = DEF_MAX_ITER)::Problem
  N::Int = length(F.ψ[1])
  S::Int = size(p,1)

  Φ::Problem = Problem(F, N, S, p, dtol, maxiter)
  findpolicy!(Φ)

  #display(Φ.F.V)
  policy::Matrix{Int} = Matrix{Int}(undef, Φ.N, Φ.S)

  for s ∈ 1:Φ.S
    policy[:,s] .= ((n)->(maximum(Φ.F.R[:,n,s]) > 0.0) ? (Φ.F.Π[n,s])[1] : -1.0).(1:Φ.N)
  end

  for n ∈ 1:Φ.N
    println(policy[n,:])
  end
  #display(policy)
  return Φ
end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
using LinearAlgebra

"""
Go from `X'Q*X + U'R*U`
s.t. x+ = Ax+Bu
to
`U'F*U + 2U'*G*x(0) + x(0)'H*x(0)`
and M,N, where X = M*U + N*x(0)
"""
function complact_form(A::AbstractArray{T}, B, Q, R, n) where T
    nu = size(B,2)
    nx = size(A,1)
    M = zeros(T, n*nx, n*nu)
    for i = 0:(n-1)
        AiB = A^(i)*B
        for j = 1:(n-i)
            rows = (nx*(i+j-1)+1):(nx*(i+j))
            cols = (nu*(j-1)+1):(nu*j)
            M[rows, cols] .= AiB
        end
    end
    N = vcat([A^i for i in 1:n]...)
    # Blockdiag
    Qn = vcat([[zeros(T, size(Q,1), size(Q,1)*i) Q zeros(T, size(Q,1), size(Q,1)*(n-i-1))]  for i = 0:(n-1)]...)
    Rn = vcat([[zeros(T, size(R,1), size(R,1)*i) R zeros(T, size(R,1), size(R,1)*(n-i-1))]  for i = 0:(n-1)]...)
    F = M'Qn*M + Rn
    G = M'Qn*N
    H = N'Qn*N
    return F, G, H, M, N
end
T = Float64
A = T.([0.999  -3.008  -0.113  -1.608  ;
     0      0.986   0.048   0       ;
     0      2.083   1.009   0       ;
     0      0.053   0.05    1       ])
B = T.([-0.080 -0.635  ;
     -0.029 -0.014  ;
     -0.868 -0.092  ;
     -0.022 -0.002  ])
Q = Matrix{T}(I,4,4)
R = Matrix{T}(I,2,2)
n = 10
F, G, H, M, N = complact_form(A,B,Q,R,n)

x0 = 0.2.*ones(4)

nx = size(A, 1)
nu = size(B, 2)

xl = -0.2*ones(T, nx)
xu = 0.2*ones(T, nx)

Xl = vcat([xl for i in 1:n]...)
Xu = vcat([xu for i in 1:n]...)
Ul = Xl - N*x0
Uu = Xu - N*x0

UC = M

using QPDAS
nU = nu*n
z = G*x0
@time qp = QuadraticProgram(zeros(T, 0, nU), zeros(T, 0), [UC;-UC], [Uu; -Ul], z, F, ϵ=1e-6, smartstart=true)

@time sol, val = solve!(qp)
val = val + 1/2*x0'H*x0

X = M*sol + N*x0
maximum(abs, X) - 0.2
maximum([UC;-UC]*sol - [Uu; -Ul])
1/2*sum(abs2, X) + 1/2*sum(abs2, sol)

max_stars_repo_name: mfalt/QPDAS

max_issues_repo_name: mfalt/QPDAS

max_forks_repo_name: mfalt/QPDAS


text:
@testset "Find Center" begin

	P = [
		0.0 1.0 0.0 0.0
		0.0 0.0 1.0 0.0
		0.0 0.0 0.0 1.0
	]

	@testset "Points Find Center" begin
		@test AlphaStructures.findCenter(P[1:1, 2:2]) == [1.0]
		@test AlphaStructures.findCenter(P[1:2, 2:2]) == [1.0, 0.0]
		@test AlphaStructures.findCenter(P[1:3, 2:2]) == [1., 0.0, 0.0]
	end

	@testset "Edges Find Center" begin
		@test AlphaStructures.findCenter(P[1:1, 1:2]) == [0.5]
		@test AlphaStructures.findCenter(P[1:2, 1:2]) == [0.5, 0.0]
		@test AlphaStructures.findCenter(P[1:3, 1:2]) == [0.5, 0.0, 0.0]
	end

	@testset "Triangles Find Center" begin
		@test AlphaStructures.findCenter(P[1:2, 1:3]) == [0.5, 0.5]
		@test AlphaStructures.findCenter(P[1:3, 1:3]) == [0.5, 0.5, 0.0]
	end

	@testset "Tetrahedrons Find Center" begin
		@test AlphaStructures.findCenter(P[1:3, 1:4]) == [0.5, 0.5, 0.5]
	end

end

#–------------------------------------------------------------------------------

@testset "Find Closest Point" begin
	P = [
		0.0 1.0 2.0 3.0 4.0
		0.0 3.0 2.0 2.0 1.0
		0.0 5.0 6.0 2.0 2.0
	]

	@testset "1D closest point" begin
		idxs = [2; 3; 4; 5]
		@test AlphaStructures.findClosestPoint(P[[1], [1]], P[[1], idxs]) == 1
		@test AlphaStructures.findClosestPoint(P[[2], [1]], P[[2], idxs]) == 4
		@test AlphaStructures.findClosestPoint(P[[3], [1]], P[[3], idxs]) == 3
	end

	@testset "2D closest point" begin
		idxs = [2; 3; 4; 5]
		σ = [1]
		while length(σ) <= 2
			newidx = idxs[
					AlphaStructures.findClosestPoint(P[1:2, σ], P[1:2, idxs])
			]
			idxs = [i for i in idxs if i ≠ newidx]
			σ = [σ; newidx]
		end
		@test σ == [1; 3; 2]
	end

	@testset "3D closest point" begin
		idxs = [2; 3; 4; 5]
		σ = [1]
		while length(σ) <= 3
			newidx = idxs[
					AlphaStructures.findClosestPoint(P[1:3, σ], P[1:3, idxs])
			]
			idxs = [i for i in idxs if i ≠ newidx]
			σ = [σ; newidx]
		end
		@test σ == [1; 4; 5; 3]
	end

	@testset "Keyword Argument Exception" begin
		@test_throws AssertionError AlphaStructures.findClosestPoint(
			P, P, metric = "Euclidean"
		)
		@test_throws AssertionError AlphaStructures.findClosestPoint(
			P, P, metric = "DD"
		)
		@test_throws AssertionError AlphaStructures.findClosestPoint(
			P, P, metric = "abc"
		)
	end

	@testset "Cannot add more points" begin
		@test_throws AssertionError AlphaStructures.findClosestPoint(P, P)
	end

end

#-------------------------------------------------------------------------------

@testset "FindMedian" begin

	P = [
		0.0 1.0 0.0 0.0
		0.0 0.0 2.0 0.0
		0.0 0.0 0.0 3.0
	]

	@test AlphaStructures.findMedian(P, 1) == 0.5
	@test AlphaStructures.findMedian(P, 2) == 1.0
	@test AlphaStructures.findMedian(P, 3) == 1.5

end

#-------------------------------------------------------------------------------

@testset "Find Radius" begin

	P = [
		0.0 1.0 0.0 0.0 0.0
		0.0 0.0 1.0 0.0 0.0
		0.0 0.0 0.0 1.0 1.0
	]

	@testset "Points Find Radius" begin
		@test AlphaStructures.findRadius(P[1:1, 2:2]) == 0.0
		@test AlphaStructures.findRadius(P[1:2, 2:2]) == 0.0
		@test AlphaStructures.findRadius(P[1:3, 2:2]) == 0.0
	end

	@testset "Edges Find Radius" begin
		@test AlphaStructures.findRadius(P[1:1, 1:2]) == 0.5
		@test AlphaStructures.findRadius(P[1:2, 1:2]) == 0.5
		@test AlphaStructures.findRadius(P[1:3, 1:2]) == 0.5

		@test AlphaStructures.findRadius(P[1:2, 4:5]) == 0.0
	end

	@testset "Triangles Find Radius" begin
		r = 0.707106781
		@test isapprox(AlphaStructures.findRadius(P[1:2, 1:3]), r, atol=1e-9)
		@test isapprox(AlphaStructures.findRadius(P[1:3, 1:3]), r, atol=1e-9)

		@test AlphaStructures.findRadius(P[1:2, 3:5]) == Inf
		@test AlphaStructures.findRadius(P[1:3, 3:5]) == Inf
	end

	@testset "Tetrahedrons Find Radius" begin
		r = 0.866025403
		@test isapprox(AlphaStructures.findRadius(P[1:3, 1:4]), r, atol=1e-9)

		@test AlphaStructures.findRadius(P[1:3, 2:5]) == Inf
	end

end

#-------------------------------------------------------------------------------

@testset "Matrix Perturbation" begin
	P = rand(Float64, 3, 100000)

	Pperturbated = AlphaStructures.matrixPerturbation(P, atol = 1e-2)
	P1 = AlphaStructures.matrixPerturbation(P, row = [1], atol = 1e-3)
	P2 = AlphaStructures.matrixPerturbation(P, col = [2], atol = 10)
	P13 = AlphaStructures.matrixPerturbation(P, row = [1; 3], atol = 1e-1)
	P11 = AlphaStructures.matrixPerturbation(P, row = [1], col = [1], atol=100)
	noPt = AlphaStructures.matrixPerturbation(P, atol = 0.0)

	@test sum(abs.(P - Pperturbated) .> 1e-2) == 0
	@test sum(abs.(P[1, :] - P1[1, :]) .> 1e-3) == 0
	@test P[[2;3], :] == P1[[2;3], :]
	@test sum(abs.(P[:, 2] - P2[:, 2]) .> 10) == 0
	@test P[:, [1;3]] == P2[:, [1;3]]
	@test sum(abs.(P[[1;3], :] - P13[[1;3], :]) .> 1e-1) == 0
	@test P[2, :] == P13[2, :]
	@test abs.(P[1, 1] - P11[1, 1]) .< 100
	@test P[2:end] == P11[2:end]
	@test P != P1 != P2 != P13 != P11 != Pperturbated
	@test P == noPt

end

#-------------------------------------------------------------------------------

@testset "Opposit Half Space" begin
	V = [
		0.0 1.0 0.0 0.0 4.0 -1. 1.0
		0.0 0.0 1.0 0.0 1.0 0.0 1.0
		0.0 0.0 0.0 1.0 2.0 0.0 1.0
	]

	@testset "1D Opposite Half Space" begin
		V1D = V[[1], :]
		@test AlphaStructures.oppositeHalfSpacePoints(
			V1D, V1D[:, [1]], V1D[:, 2]
		) == [6]
		@test AlphaStructures.oppositeHalfSpacePoints(
			V1D, V1D[:, [6]], V1D[:, 1]
		) == []
		@test AlphaStructures.oppositeHalfSpacePoints(
			V1D, V1D[:, [1]], V1D[:, 6]
		) == [2; 5; 7]
	end

	@testset "2D Opposite Half Space" begin
		V2D = V[1:2, :]
		@test AlphaStructures.oppositeHalfSpacePoints(
			V2D, V2D[:, [2; 3]], V2D[:, 1]
		) == [5; 7]
		@test AlphaStructures.oppositeHalfSpacePoints(
			V2D, V2D[:, [1; 2]], V2D[:, 3]
		) == []
		@test AlphaStructures.oppositeHalfSpacePoints(
			V2D, V2D[:, [1; 3]], V2D[:, 2]
		) == [6]
	end

	@testset "3D Opposite Half Space" begin
		@test AlphaStructures.oppositeHalfSpacePoints(
			V, V[:, [2; 3; 4]], V[:, 1]
		) == [5; 7]
		@test AlphaStructures.oppositeHalfSpacePoints(
			V, V[:, [1; 2; 3]], V[:, 4]
		) == []
		@test AlphaStructures.oppositeHalfSpacePoints(
			V, V[:, [1; 3; 4]], V[:, 2]
		) == [6]
	end

end

#-------------------------------------------------------------------------------

@testset "planarIntersection" begin

	P = [
		0.0 1.0 -1. 1.0 2.0
		1.0 2.0 -1. 1.0 3.0
		1.0 0.0 -1. -1. -5.
	]

	@testset "1D Planar Intersection" begin
		@test AlphaStructures.planarIntersection(P, [1], 1, 0.0) == 0
		@test AlphaStructures.planarIntersection(P, [2], 1, 0.0) == 1
		@test AlphaStructures.planarIntersection(P, [3], 1, 0.0) == -1
	end

	@testset "2D Planar Intersection" begin
		@test AlphaStructures.planarIntersection(P, [2; 3], 1, 0.0) == 0
		@test AlphaStructures.planarIntersection(P, [1; 2], 2, 0.0) == 1
		@test AlphaStructures.planarIntersection(P, [2; 3], 3, 0.0) == -1
	end

	@testset "3D Planar Intersection" begin
		@test AlphaStructures.planarIntersection(P, [1; 2; 3], 1, 0.0) == 0
		@test AlphaStructures.planarIntersection(P, [1; 2; 4], 2, 0.0) == 1
		@test AlphaStructures.planarIntersection(P, [2; 3; 4], 3, 0.0) == -1
	end

	@testset "4D Planar Intersection" begin
		@test AlphaStructures.planarIntersection(P, [2; 3; 4; 5], 1, 0.0) == 0
		@test AlphaStructures.planarIntersection(P, [1; 2; 4; 5], 2, 0.0) == 1
		@test AlphaStructures.planarIntersection(P, [2; 3; 4; 5], 3, 0.0) == -1
	end
end

#-------------------------------------------------------------------------------

@testset "Simplex Faces" begin
	@test AlphaStructures.simplexFaces([1; 2]) == [ [1], [2] ]
	@test AlphaStructures.simplexFaces([1; 2; 3]) ==
		[ [1, 2], [1, 3], [2, 3] ]
	@test AlphaStructures.simplexFaces([1; 2; 3; 4]) ==
		[ [1, 2, 3], [1, 2, 4], [1, 3, 4], [2, 3, 4] ]
	@test AlphaStructures.simplexFaces([4; 3; 2; 1; 5]) ==
		[ [1, 2, 3, 4], [1, 2, 3, 5], [1, 2, 4, 5], [1, 3, 4, 5], [2, 3, 4, 5] ]
end

#-------------------------------------------------------------------------------

 @testset "Vertex in Circumball" begin

 	@testset "2D Vertex in Circumball" begin
 		V = [
 			0. 1. 0.;
 			0. 0. 1.
 		]
 		simplex = [2, 3]
 		up_simplex = [1, 2, 3]
 		point = V[:, setdiff(up_simplex, simplex)]
 		T= V[:, simplex]
 		@test AlphaStructures.vertexInCircumball(
 			T, AlphaStructures.findRadius(T), point
 		)
 	end

 	@testset "3D Vertex in Circumball" begin

 		@testset "edge and triangle" begin
 			V=[
 				0. 1. 0.;
 				0. 0. 1.;
 				0. 0. 0.
 			]
 			simplex = [2, 3]
 			up_simplex = [1, 2, 3]
 			point = V[:, setdiff(up_simplex, simplex)]
 			T= V[:, simplex]
			@test AlphaStructures.vertexInCircumball(
 				T, AlphaStructures.findRadius(T), point
 			)
 		end

 		@testset "triangle and tetrahedron" begin
 			V=[
 				0. 1. 0. 0.;
 				0. 0. 1. 0.;
 				0. 0. 0. 1.
 			]
 			simplex = [2, 3, 4]
 			up_simplex = [1, 2, 3, 4]
 			point = V[:, setdiff(up_simplex, simplex)]
 			T= V[:, simplex]
 			@test AlphaStructures.vertexInCircumball(
 				T, AlphaStructures.findRadius(T), point
 			)
 		end

 	end

 end


text:
#!/usr/bin/env bash
#=
export JULIA_DEBUG=UnivariateExperiments
exec julia --project="." --color=yes --startup-file=no -e \
    'include(popfirst!(ARGS))' "${BASH_SOURCE[0]}" "$@"
=#

using JuMP
using CPLEX
using PolyhedralRelaxations
using Printf
using Plots
using SparseArrays
using ArgParse
using DelimitedFiles

include("src/basics.jl")
include("src/plotting.jl")
include("src/minlplib_helper.jl")
include("src/trig.jl")
include("src/ex4_1_1.jl")
include("src/fo7.jl")
include("src/errorf.jl")
include("src/ramsey.jl")
include("src/gammaf.jl")

s = ArgParseSettings()
@add_arg_table! s begin
    "--plot", "-p"
    help = "generate plots"
    action = :store_true
    "--minlplib", "-m"
    help = "generate results for MINLPLIB instances"
    action = :store_false
    "--basic", "-b"
    help = "collect basic results for atomic functions"
    action = :store_true
end

parsed_args = parse_args(ARGS, s)

if parsed_args["plot"]
    generate_plots()
end

if parsed_args["minlplib"]
    # @info "running the relaxations on trig"
    # trig()
    # @info "running the relaxations on ex4_1_1"
    # ex4_1_1()
    # @info "running the relaxations on fo7"
    # fo7()
    @info "running the relaxations on errorf"
    errorf()
    @info "running the relaxations on ramsey"
    ramsey()
    @info "running the relaxations on gamma"
    gammaf()
end

if parsed_args["basic"]
    @info "generating error result table..."
    generate_error_table()
    @info "generated error resul table."
end

max_stars_repo_name: sujeevraja/univariate-experiments

max_issues_repo_name: sujeevraja/univariate-experiments

max_forks_repo_name: sujeevraja/univariate-experiments


text:
gpuvec(x) = GPUVector(GLBuffer(x))

function to_nd(x, n::Type{Val{N}}, default) where N
    ntuple(n) do i
        i <= length(x) && return x[i]
        default
    end
end

function to_world(point::T, cam) where T <: StaticVector
    x = to_world(
        point,
         inv(Reactive.value(cam.view)) * inv(Reactive.value(cam.projection)),
        T(widths(Reactive.value(cam.window_size)))
    )
    Point2f0(x[1], x[2])
end
function to_world(
        p::StaticVector{N, T},
        prj_view_inv::Mat4,
        cam_res::StaticVector
    ) where {N, T}
    VT = typeof(p)
    clip_space = ((VT(p) ./ VT(cam_res)) .* T(2)) .- T(1)
    pix_space = Vec{4, T}(
        clip_space[1],
        clip_space[2],
        T(0), GLAbstraction.w_component(p)
    )
    ws = prj_view_inv * pix_space
    ws ./ ws[4]
end

function qrotation(axis::StaticVector{N, T}, theta) where {N, T <: Real}
    if length(axis) != 3
        error("Must be a 3-vector")
    end
    u = normalize(axis)
    thetaT = convert(eltype(u), theta)
    s = sin(thetaT / T(2))
    Vec4f0(s * u[1], s * u[2], s * u[3], cos(thetaT / T(2)))
end

qabs(q) = sqrt(dot(q, q))

function qnormalize(q)
    q ./ qabs(q)
end

function qmul(quat::StaticVector{4}, vec::StaticVector{2})
    x3 = qmul(quat, Vec(vec[1], vec[2], 0))
    StaticArrays.similar_type(vec, StaticArrays.Size(2,))(x3[1], x3[2])
end
function qmul(quat::StaticVector{4}, vec::StaticVector{3})
    num = quat[1] * 2f0;
    num2 = quat[2] * 2f0;
    num3 = quat[3] * 2f0;
    num4 = quat[1] * num;
    num5 = quat[2] * num2;
    num6 = quat[3] * num3;
    num7 = quat[1] * num2;
    num8 = quat[1] * num3;
    num9 = quat[2] * num3;
    num10 = quat[4] * num;
    num11 = quat[4] * num2;
    num12 = quat[4] * num3;
    return Point3f0(
        (1f0 - (num5 + num6)) * vec[1] + (num7 - num12) * vec[2] + (num8 + num11) * vec[3],
        (num7 + num12) * vec[1] + (1f0 - (num4 + num6)) * vec[2] + (num9 - num10) * vec[3],
        (num8 - num11) * vec[1] + (num9 + num10) * vec[2] + (1f0 - (num4 + num5)) * vec[3]
    )
end
qconj(q) = Vec4f0(-q[1], -q[2], -q[3], q[4])

function qmul(q::StaticVector{4}, w::StaticVector{4})
    qq = Vec4f0(
        q[4] * w[4] - q[1] * w[1] - q[2] * w[2] - q[3] * w[3],
        q[4] * w[1] + q[1] * w[4] + q[2] * w[3] - q[3] * w[2],
        q[4] * w[2] - q[1] * w[3] + q[2] * w[4] + q[3] * w[1],
        q[4] * w[3] + q[1] * w[2] - q[2] * w[1] + q[3] * w[4],
    )
    qnormalize(qq)
end

is_unitrange(x) = (false, 0:0)
is_unitrange(x::Range) = (true, x)
function is_unitrange(x::AbstractVector)
    length(x) < 2 && return false, 0:0
    diff = x[2] - x[1]
    length(x) < 3 && return true, x[1]:x[2]
    last = x[3]
    for elem in drop(x, 3)
        diff2 = elem - last
        diff2 != diff && return false, 0:0
    end
    return true, range(first(x), diff, length(x))
end

function ngrid(x::AbstractVector, y::AbstractVector)
    xgrid = [Float32(x[i]) for i = 1:length(x), j = 1:length(y)]
    ygrid = [Float32(y[j]) for i = 1:length(x), j = 1:length(y)]
    xgrid, ygrid
end

function nan_extrema(array)
    mini, maxi = (Inf, -Inf)
    for elem in array
        isnan(elem) && continue
        mini = min(mini, elem)
        maxi = max(maxi, elem)
    end
    Vec2f0(mini, maxi)
end