text:
module DPO
export rewrite, rewrite_match, valid_dpo, dangling_condition, id_condition,
  pushout_complement

using ..FinSets, ..CSets, ..Limits, ..Subobjects
using ...Theories
using ...Theories: attr

"""
    l
  L ← I
m ↓   ↓k
  G ← K
    g

Compute a pushout complement, componentwise in Set. Formally, for each
component, define K = G / m(L/l(I)). There is a natural injection g: K↪G. For
each component, define k as equal to the map l;m (every element in the image in
G is also in K).

Returns ACSetTransformations k and g such that (m, g) is the pushout of (l, k).
Elements of K are ordered in the same order as they appear in G.
"""
function pushout_complement(
    l::ACSetTransformation{CD,AD}, m::ACSetTransformation{CD,AD}
  )::Pair{ACSetTransformation{CD,AD},ACSetTransformation{CD,AD}} where {CD,AD}
  valid_dpo(l, m) || error("morphisms L and m do not satisfy gluing conditions")
  I, L, G = dom(l), codom(l), codom(m)

  # Construct subobject g: K ↪ G.
  g_components = NamedTuple{ob(CD)}(map(ob(CD)) do c
    l_image = Set(collect(l[c]))
    orphans = Set([ m[c](x) for x in parts(L,c) if x ∉ l_image ])
    filter(x -> x ∉ orphans, parts(G,c))
  end)
  g = hom(Subobject(G, g_components))
  K = dom(g)

  # Construct morphism k: I → K using partial inverse of g.
  lm = compose(l, m)
  k_components = NamedTuple{ob(CD)}(map(ob(CD)) do c
    g_inv = Dict{Int,Int}(zip(g_components[c], parts(K,c)))
    [ g_inv[lm[c](x)] for x in parts(I,c) ]
  end)
  k = ACSetTransformation(k_components, I, K)

  return k => g
end

"""
Apply a rewrite rule (given as a span, L<-I->R) to a ACSet
using a match morphism `m` which indicates where to apply
the rewrite.
"""
function rewrite_match(L::ACSetTransformation{CD, AD},
                       R::ACSetTransformation{CD, AD},
                       m::ACSetTransformation{CD, AD}
                      )::AbstractACSet{CD, AD} where {CD, AD}
    dom(L) == dom(R) || error("Rewriting where L, R do not share domain")
    codom(L) == dom(m) || error("Rewriting where L does not compose with m")
    (k, _) = pushout_complement(L, m)
    l1, _ = pushout(R, k)
    return codom(l1)
end

"""
Apply a rewrite rule (given as a span, L<-I->R) to a ACSet,
where a match morphism is found automatically. If multiple
matches are found, a particular one can be selected using
`m_index`.
"""
function rewrite(L::ACSetTransformation{CD, AD},
                 R::ACSetTransformation{CD, AD},
                 G::AbstractACSet{CD, AD},
                 monic::Bool=false, m_index::Int=1
                )::Union{Nothing, AbstractACSet} where {CD, AD}
  ms = filter(m->valid_dpo(L, m), homomorphisms(codom(L), G, monic=monic))
  if 0 < m_index <= length(ms)
    return rewrite_match(L, R, ms[m_index])
  else
    return nothing
  end
end

"""
Condition for existence of a pushout complement
"""
function valid_dpo(L::ACSetTransformation, m::ACSetTransformation)::Bool
  return all(isempty, [collect(id_condition(L, m))...,
                       dangling_condition(L, m)])
end

"""
Check the identification condition: Does not map both a deleted item and a
preserved item in L to the same item in G, or two distinct deleted items to the
same. (Trivially satisfied if mono)

Returns a tuple of lists of violations
  1.) For a given component, a pair of IDs in L that are deleted yet mapped to
      the same index (the last integer of the tuple) in G
  2.) For a given component, a nondeleted index that maps to a deleted index
      in G
"""
function id_condition(L::ACSetTransformation{CD, AD},
                      m::ACSetTransformation{CD, AD}) where {CD, AD}
  res1, res2 = Tuple{Symbol, Int, Int, Int}[], Tuple{Symbol, Int, Int}[]
  for comp in keys(components(L))
    m_comp = x->m[comp](x)
    image = Set(collect(L[comp]))
    image_complement = filter(x->!(x in image), parts(codom(L),comp))
    image_vals = map(m_comp, collect(image))
    orphan_vals = map(m_comp, image_complement)
    orphan_set = Set(orphan_vals)

    for (i, iv) in enumerate(image_complement)
      for j in i+1:length(image_complement)
        if m_comp(iv) == m_comp(image_complement[j])
          push!(res1, (comp, i, j, m_comp(i)))
        end
      end
    end
    for i in image
      if m_comp(i) in orphan_set
        push!(res2, (comp, i, m_comp(i)))
      end
    end
  end

  return (res1, res2)
end

"""
Check the dangling condition: m doesn't map a deleted element d to a element
m(d) ∈ G if m(d) is connected to something outside the image of m.

For example, in the CSet of graphs:
  e1
1 --> 2

if e1 is not matched but either 1 and 2 are deleted, then e1 is dangling
"""
function dangling_condition(L::ACSetTransformation{CD, AD},
                            m::ACSetTransformation{CD, AD}) where {CD, AD}
  orphans, res = Dict(), []
  for comp in keys(components(L))
    image = Set(collect(L[comp]))
    orphans[comp] = Set(
      map(x->m[comp](x),
        filter(x->!(x in image),
          parts(codom(L), comp))))
  end
  # check that for all morphisms in C, we do not map to an orphan
  for (morph, src_ind, tgt_ind) in zip(hom(CD), dom(CD), codom(CD))
    src_obj = ob(CD)[src_ind] # e.g. :E, given morph=:src in graphs
    tgt_obj = ob(CD)[tgt_ind] # e.g. :V, given morph=:src in graphs
    n_src = parts(codom(m), src_obj)
    unmatched_vals = setdiff(n_src, collect(m[src_obj]))
    unmatched_tgt = map(x -> m.codom[morph][x], collect(unmatched_vals))
    for unmatched_val in setdiff(n_src, collect(m[src_obj]))  # G/m(L) src
      unmatched_tgt = m.codom[morph][unmatched_val]
      if unmatched_tgt in orphans[tgt_obj]
        push!(res, (morph, unmatched_val, unmatched_tgt))
      end
    end
  end
  return res
end

end


text:
using Documenter

makedocs(
    format = :html,
    sitename = "AOTDemo.jl",
    pages = [
        "index.md",
    ]
)

deploydocs(
    repo = "github.com/tangentworks/AOTDemo.jl.git",
    target = "build",
    julia  = "0.6",
    deps = nothing,
    make = nothing,
    osname = "linux"
)

ext: jl


text:
module ImageProcessing

using FileIO
using ColorTypes
using DeformableRegistration.Types
using DeformableRegistration.Transformation
using DeformableRegistration.Interpolation
using MicroLogging
export createImage, loadImage, setImageProperties
export restrictResolutionToLevel
export getSize, getSpatialDomain, getPixelSpacing

include("./helpers/smoothing.jl")
export smoothArray

# load image and convert it into gray image
function loadImage(pathToFile;
                   voxelsize = [1.0,1.0], shift = [0.0, 0.0])

    # load image
    image = load(pathToFile)
    #image = permutedims(image, (2,1))

    image = convert(Array{ColorTypes.Gray},image)

    # convert image data to float (easier access)
    image = convert(Array{Float64,2},image)

    # set image properties

    return regImage(image, voxelsize, shift)

end

function createImage(imageData;
                     voxelsize = [1.0,1.0], shift = [0.0, 0.0])

    # load image
    #image = grayim(copy(imageData))

    return regImage(imageData, voxelsize, shift)

end

function restrictRegImage(im::regImage)
    newImage = deepcopy(im)
    newImage.data[:] = smoothArray(newImage.data, 3, 3.0)
    oldsize = size(im)
    newsize = map(Int,ceil.(oldsize./2))
    newGrid = getCellCenteredGrid(im.voxelsize*2, im.shift, newsize)
    newImage = createImage(interpolateImage(newImage,newGrid)[1], voxelsize = newGrid.voxelsize, shift = newGrid.shift)
    return newImage
end

function restrictResolutionToLevel(im::regImage,level)

    # copy restrict image, there seems to be a bug in the image copy method, take care of the deepcopy here
    #im = deepcopy(image.data)
    maxlevel = 0;
    for l=1:level
        if( (size(im.data,1)>2) && (size(im.data,2)>2) )
            im = restrictRegImage(im)
        end
    end
    return im
end

function getSize(I::regImage)
  return [size(I.data,1), size(I.data,2)]
end

end


text:
using DiffEqFlux, Flux, Optim, OrdinaryDiffEq

n = 5  # number of ODEs
tspan = (0.0, 1.0)

d = 5  # number of data pairs
x_data = rand(n, 5)
y_data = rand(n, 5)
training_data = Iterators.cycle([(x_data[:, i], y_data[:, i]) for i in 1:d])

NN = Chain(Dense(n, 10n, tanh),
           Dense(10n, n))

@info "ROCK4"
nODE = NeuralODE(NN, tspan, ROCK4(), reltol=1e-4, saveat=[tspan[end]])

loss_function(θ, x, y) = Flux.mse(y, nODE(x, θ))
res = DiffEqFlux.sciml_train(loss_function, nODE.p, LBFGS(), training_data)
res = DiffEqFlux.sciml_train(loss_function, nODE.p, NewtonTrustRegion(), training_data)

NN = SlowChain(SlowDense(n, 10n, tanh),
               SlowDense(10n, n))

@info "ROCK2"
nODE = NeuralODE(NN, tspan, ROCK2(), reltol=1e-4, saveat=[tspan[end]])

loss_function(θ, x, y) = Flux.mse(y, nODE(x, θ))
res = DiffEqFlux.sciml_train(loss_function, nODE.p, LBFGS(), training_data)
res = DiffEqFlux.sciml_train(loss_function, nODE.p, NewtonTrustRegion(), training_data)

@info "ROCK4"
nODE = NeuralODE(NN, tspan, ROCK4(), reltol=1e-4, saveat=[tspan[end]])

loss_function(θ, x, y) = Flux.mse(y, nODE(x, θ))
res = DiffEqFlux.sciml_train(loss_function, nODE.p, LBFGS(), training_data)
res = DiffEqFlux.sciml_train(loss_function, nODE.p, NewtonTrustRegion(), training_data)


text:
module CalibrateEDMF

# Submodules
include("HelperFuncs.jl")
include("ModelTypes.jl")
include("DistributionUtils.jl")
include("ReferenceModels.jl")
include("LESUtils.jl")
include("ReferenceStats.jl")
include("KalmanProcessUtils.jl")
include("TurbulenceConvectionUtils.jl")
include("Diagnostics.jl")
include("NetCDFIO.jl")
include("Pipeline.jl")

end # module


text:
export
    RBState,
    randbetween

"""
    RBState{T} <: StaticVector{13,T}

Represents the state of a rigid body in 3D space, consisting of position, orientation, linear
    velocity and angular velocity, respresented as a vector stacked in that order, with
    the rotation represented as the 4 elements of a unit quaternion.

Implements the `StaticArrays` interface so can be treated as an `SVector` with additional
    methods.

# Constructors
    RBState{T}(r, q, v, ω)
    RBState{T}(x)
    RBState(r, q, v, ω)
    RBState(x)

where `r`, `v`, and `ω` are three-dimensional vectors, `q` is either a `Rotation` or a
    four-dimenional vector representing the parameters of unit quaternion, and `x` is a
    13-dimensional vector (or tuple),

# Converting to a State Vector
An `RBState` can be converted to a state vector for a `RigidBody` using
    RBState(model::RBstate, x, [renorm=false])
"""
struct RBState{T} <: StaticVector{13,T}
    r::SVector{3,T}
    q::UnitQuaternion{T}
    v::SVector{3,T}
    ω::SVector{3,T}
    function RBState{T}(r, q::Rotation, v, ω) where T
        @assert length(r) == 3
        @assert length(v) == 3
        @assert length(ω) == 3
        new{T}(r, UnitQuaternion{T}(q), v, ω)
    end
    @inline function RBState{T}(x::RBState) where T
        RBState{T}(x.r, x.q, x.v, x.ω)
    end
end

function RBState(r::AbstractVector, q::Rotation, v::AbstractVector, ω::AbstractVector)
    T = promote_type(eltype(r), eltype(q), eltype(v), eltype(ω))
    RBState{T}(r, q, v, ω)
end

@inline function RBState(r::AbstractVector, q::AbstractVector, v::AbstractVector, ω::AbstractVector)
    RBState(r, UnitQuaternion(q, false), v, ω)
end

@inline RBState(x::RBState) = x

@generated function RBState(::Type{R}, x::AbstractVector) where R<:Rotation{3}
    ir = SA[1,2,3]
    iq = SA[4,5,6]
    iv = SA[7,8,9]
    iω = SA[10,11,12]
    if R <: UnitQuaternion
        iq = SA[4,5,6,7]
        iv = iv .+ 1
        iω = iω .+ 1
    end
    quote
        q = UnitQuaternion(R(x[$iq]))
        RBState(x[$ir], q, x[$iv], x[$iω])
    end
end

# Static Arrays interface
function (::Type{RB})(x::NTuple{13}) where RB <: RBState
    RB(
        SA[x[1], x[2], x[3]],
        UnitQuaternion(x[4], x[5], x[6], x[7], false),
        SA[x[8], x[9], x[10]],
        SA[x[11], x[12], x[13]]
    )
end
Base.@propagate_inbounds function Base.getindex(x::RBState, i::Int)
    if i < 4
        x.r[i]
    elseif i < 8
        Rotations.params(x.q)[i-3]
    elseif i < 11
        x.v[i-7]
    else
        x.ω[i-10]
    end
end
Base.Tuple(x::RBState) = (
    x.r[1], x.r[2], x.r[3],
    x.q.w, x.q.x, x.q.y, x.q.z,
    x.v[1], x.v[2], x.v[3],
    x.ω[1], x.ω[2], x.ω[3]
)

"""
    renorm(x::RBState)

Re-normalize the unit quaternion.
"""
@inline renorm(x::RBState) = RBState(x.r, normalize(x.q), x.v, x.ω)

"""
    position(x::RBState)
    position(model::RigidBody, x::AbstractVector)

Return the 3-dimensional position of an rigid body as a `SVector{3}`.
"""
@inline Base.position(x::RBState) = x.r

"""
    orientation(x::RBState)
    orientation(model::RigidBody, x::AbstractVector)

Return the 3D orientation of a rigid body. Returns a `Rotations.Rotation{3}`.
"""
@inline orientation(x::RBState) = x.q

"""
    angular_velocity(x::RBState)
    angular_velocity(model::RigidBody, x::AbstractVector)

Return the 3D linear velocity of a rigid body as a `SVector{3}`.
"""
@inline angular_velocity(x::RBState) = x.ω

"""
    linear_velocity(x::RBState)
    linear_velocity(model::RigidBody, x::AbstractVector)

Return the 3D linear velocity of a rigid body as a `SVector{3}`.
"""
@inline linear_velocity(x::RBState) = x.v

function Base.isapprox(x1::RBState, x2::RBState; kwargs...)
    isapprox(x1.r, x2.r; kwargs...) &&
        isapprox(x1.v, x2.v; kwargs...) &&
        isapprox(x1.ω, x2.ω; kwargs...) &&
        isapprox(principal_value(x1.q), principal_value(x2.q); kwargs...)
end

"""
    +(::RBState, ::RBState)

Add two rigid body states, which adds the position, linear and angular velocities, and
    composes the orientations.
"""
function Base.:+(s1::RBState, s2::RBState)
    RBState(s1.r+s2.r, s1.q*s2.q, s1.v+s2.v, s1.ω+s2.ω)
end

"""
    -(::RBState, ::RBState)

Substract two rigid body states, which substracts the position, linear and angular velocities,
    and composes the inverse of the second orientation with the first, i.e. `inv(q2)*q1`.
"""
function Base.:-(s1::RBState, s2::RBState)
    RBState(s1.r-s2.r, s2.q\s1.q, s1.v-s2.v, s1.ω-s2.ω)
end

"""
    ⊖(::RBState, ::RBState)

Compute the 12-dimensional error state, calculated by substracting thep position, linear,
    and angular velocities, and using `Rotations.rotation_error` for the orientation.
"""
function Rotations.:⊖(s1::RBState, s2::RBState, rmap=CayleyMap())
    dx = s1.r-s2.r
    dq = Rotations.rotation_error(s1.q, s2.q, rmap)
    dv = s1.v-s2.v
    dw = s1.ω-s2.ω
    @SVector [dx[1], dx[2], dx[3], dq[1], dq[2], dq[3],
              dv[1], dv[2], dv[3], dw[1], dw[2], dw[3]]
end

"""
    ⊕(x::RBState, dx::StaticVector{12}, rmap=CayleyMap())

Add the state error `dx` to the rigid body state `x` using the map `rmap`. Simply adds the
vector states, and computes the orientation with `Rotations.add_error(x, dx, rmap)`
"""
function Rotations.:⊕(s1::RBState, dx::StaticVector{12}, rmap=CayleyMap())
    dr = SA[dx[1],  dx[2],  dx[3]]
    dq = SA[dx[4],  dx[5],  dx[6]]
    dv = SA[dx[7],  dx[8],  dx[9]]
    dω = SA[dx[10], dx[11], dx[12]]
    q = Rotations.add_error(s1.q, Rotations.RotationError(dq, rmap))
    RBState(s1.r + dr, q, s1.v + dv, s1.ω + dω)
end

Base.zero(s1::RBState) = zero(RBState)
@inline Base.zero(::Type{<:RBState}) = zero(RBState{Float64})
function Base.zero(::Type{<:RBState{T}}) where T
    RBState((@SVector zeros(T,3)), UnitQuaternion{T}(I),
        (@SVector zeros(T,3)), (@SVector zeros(T,3)))
end

function Base.rand(::Type{RBState})
    RBState(rand(3), rand(UnitQuaternion), rand(3), rand(3))
end

function randbetween(xmin::RBState, xmax::RBState)
    dx = xmax - xmin
    RBState(
        xmin.r .+ rand(3) .* dx.r,
        # rand(UnitQuaternion),
        expm((@SVector randn(3))*rand()*deg2rad(170)),
        xmin.v .+ rand(3) .* dx.v,
        xmin.ω .+ rand(3) .* dx.ω
    )
end


text:
using PIPG
import MathOptInterface
using LinearAlgebra
using StaticArrays
using SparseArrays
using Profile
const MOI = MathOptInterface
using Test

@testset "Linear algebra routines" begin
    for i=1:100
        res = MVector{10,Float64}(zeros(10))
        mat = sprand(10,10,0.5)
        inp = SVector{10, Float64}(rand(10))
        PIPG.spmul!(res, mat, inp, 1.0, 0.0)
        @test all(abs.(res .- mat*inp) .< 1e-10)
        PIPG.spmul!(res, mat, inp, 1.0, -1.0)
        @test all(abs.(res) .< 1e-10)
        PIPG.spmul!(res, mat, inp, -1.0, 0.0)
        @test all(abs.(res .+ mat*inp) .< 1e-10)
    end
end

@testset "Basic optimizer examples" begin
	γ = 0.9
	# QP example
    function make_qp()
    	P_ex = spdiagm([1/MathConstants.e, 2.0])
    	q = MVector(1/MathConstants.e,2.0)
    	H_ex = sparse([-1.0 -1.0])
    	g = MVector(0.0)
    	return PIPG.Problem(PIPG.NOCone{Float64, 1}(), PIPG.Reals{Float64, 2}(), H_ex, P_ex, q, g, 0.0)
    end

    # no scaling
    prob = make_qp()
    state = PIPG.State(prob; scaling=())
	PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.0001, SVector(1.0, 1.0), SVector(0.0))
    result = state.primal
	@test norm(result .- [0.689, -0.689]) < 0.001

	# LP example
    make_prob() = PIPG.Problem(PIPG.PTCone{Float64}((PIPG.POCone{Float64, 1}(), PIPG.POCone{Float64, 1}(), PIPG.POCone{Float64, 1}(),
        PIPG.POCone{Float64, 2}())), PIPG.Reals{Float64, 2}(),
        sparse([-100.0 -200.0; -10.0 -30.0; -1.0 -1.0; 1.0 0.0; 0.0 1.0]),
        spzeros(2,2),
        -1 .* MVector(50.0, 120.0),
        MVector(-10000.0, -1200.0, -110, -0.0, -0.0), 0.0)
	prob = make_prob()
	state = PIPG.State(prob; scaling=())
	PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.0001, SVector(0.0, 0.0), SVector(0.0, 0.0, 0.0, 0.0, 0.0))
	@test norm(state.primal .- [60, 20]) < 0.01

    # geometric scaling
    prob = make_prob()
    state = PIPG.State(prob; scaling=(PIPG.GeoMean(prob), ))
    PIPG.scale(prob, state)
    niters = PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.0001, SVector(0.0, 0.0), SVector(0.0, 0.0, 0.0, 0.0, 0.0))
    @test norm(state.primal .* state.col_scale .- [60, 20]) < 0.01

    prob = make_qp()
    state = PIPG.State(prob; scaling=(PIPG.GeoMean(prob), ))
    PIPG.scale(prob, state)
    niters = PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.0001, SVector(0.0, 0.0), SVector(0.0))
    @test norm(state.primal .* state.col_scale .- [0.689, -0.689]) < 0.001

    # arithmetic scaling
    prob = make_prob()
    state = PIPG.State(prob; scaling=(PIPG.ArithMean(prob), ))
    PIPG.scale(prob, state)
    niters = PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.0001, SVector(0.0, 0.0), SVector(0.0, 0.0, 0.0, 0.0, 0.0))
    @test norm(state.primal .* state.col_scale .- [60, 20]) < 0.01

    prob = make_qp()
    state = PIPG.State(prob; scaling=(PIPG.ArithMean(prob), ))
    PIPG.scale(prob, state)
    niters = PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.00001, SVector(0.0, 0.0), SVector(0.0))
    @test norm(state.primal .* state.col_scale .- [0.689, -0.689]) < 0.001

    # equilibration scaling
    prob = make_prob()
    state = PIPG.State(prob; scaling=(PIPG.Equilibration(prob), ))
    PIPG.scale(prob, state)
    niters = PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.0001, SVector(0.0, 0.0), SVector(0.0, 0.0, 0.0, 0.0, 0.0))
    @test norm(state.primal .* state.col_scale .- [60, 20]) < 0.01

    prob = make_qp()
    state = PIPG.State(prob; scaling=(PIPG.Equilibration(prob), ))
    PIPG.scale(prob, state)
    niters = PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.00001, SVector(0.0, 0.0), SVector(0.0))
    @test norm(state.primal .* state.col_scale .- [0.689, -0.689]) < 0.001

    # default scaling (geomean + equlibration)
    prob = make_prob()
    state = PIPG.State(prob)
    PIPG.scale(prob, state)
    niters = PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.0001, SVector(0.0, 0.0), SVector(0.0, 0.0, 0.0, 0.0, 0.0))
    @test norm(state.primal .* state.col_scale .- [60, 20]) < 0.01

    prob = make_qp()
    state = PIPG.State(prob)
    PIPG.scale(prob, state)
    niters = PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.00001, SVector(0.0, 0.0), SVector(0.0))
    @test norm(state.primal .* state.col_scale .- [0.689, -0.689]) < 0.001
end

#=
@testset "Profiling" begin
    make_prob() = PIPG.Problem(PIPG.PTCone{Float64}((PIPG.POCone{Float64, 1}(), PIPG.POCone{Float64, 1}(), PIPG.POCone{Float64, 1}(),
        PIPG.POCone{Float64, 2}())), PIPG.Reals{Float64, 2}(),
        sparse([-100.0 -200.0; -10.0 -30.0; -1.0 -1.0; 1.0 0.0; 0.0 1.0]),
        spzeros(2,2),
        -1 .* MVector(50.0, 120.0),
        MVector(-10000.0, -1200.0, -110, -0.0, -0.0), 0.0)
    prob = make_prob()
    state = PIPG.State(prob)
    γ = 0.9
    PIPG.scale(prob, state)
    scale_a_bunch(prob,state,niters) = for i=1:niters PIPG.scale(prob, state) end
    scale_a_bunch(prob, state, 1)
    Profile.clear()
    @profile scale_a_bunch(prob, state, 1000000)
    @time scale_a_bunch(prob, state, 1000000)
    Profile.print(C=true)
end
=#

function basic_lp()
    opt = PIPG.Optimizer(ϵ=1e-7)
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    c = [50.0, 120.0]
    x = MOI.add_variables(model, length(c))
    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}(),
        MOI.ScalarQuadraticFunction(MOI.ScalarQuadraticTerm{Float64}[], MOI.ScalarAffineTerm.(c, x), 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([100.0, 200.0], x), 0.0),
        MOI.LessThan(10000.0))
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([10.0, 30.0], x), 0.0),
        MOI.LessThan(1200.0))
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 1.0], x), 0.0),
        MOI.LessThan(110.0))
    MOI.add_constraint(model, MOI.VectorOfVariables(x), MOI.Nonnegatives(2))
    return opt, model, x
end

@testset "Basic LP" begin
    opt,model,x = basic_lp()
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test norm(MOI.get.(model, MOI.VariablePrimal(), x)  .- [60, 20]) < 0.01
    @test MOI.get(model, MOI.ObjectiveValue()) ≈ 5400

    opt,model,x = basic_lp()
    opt.scaling = (p->(PIPG.GeoMean(p), ))
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test norm(MOI.get.(model, MOI.VariablePrimal(), x)  .- [60, 20]) < 0.01
    @test MOI.get(model, MOI.ObjectiveValue()) ≈ 5400
end

@testset "Basic LP - VectorAffineFunction" begin
    opt = PIPG.Optimizer()
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    c = [50.0, 120.0]
    x = MOI.add_variables(model, length(c))
    o = MOI.set(model,
    	MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}(),
    	MOI.ScalarQuadraticFunction(MOI.ScalarQuadraticTerm{Float64}[], MOI.ScalarAffineTerm.(c, x), 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    cv = MOI.add_constraint(model,
    	MOI.VectorAffineFunction(MOI.VectorAffineTerm.(
    	[1,1,2,2,3,3],
    	[MOI.ScalarAffineTerm.([100.0, 200.0], x);
    	 MOI.ScalarAffineTerm.([10.0, 30.0], x);
    	 MOI.ScalarAffineTerm.([1.0, 1.0], x)]), [-10000.0, -1200.0, -110.0]),
    	MOI.Nonpositives(3))
    pos = MOI.add_constraint(model, MOI.VectorOfVariables(x), MOI.Nonnegatives(2))
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test norm(MOI.get.(model, MOI.VariablePrimal(), x)  .- [60, 20]) < 0.01
    @test MOI.get(model, MOI.ObjectiveValue()) ≈ 5400
end

@testset "Basic SOCP" begin
    opt,model,x = basic_lp()
    opt.scaling = p->(PIPG.ArithMean(p),)
    tip = MOI.add_variable(model)
    MOI.add_constraint(model, tip, MOI.EqualTo(50.0))
    MOI.add_constraint(model, MOI.VectorOfVariables([tip, x...]),
        MOI.SecondOrderCone(3)) # |x|_2 < 50.0
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test norm(MOI.get.(model, MOI.VariablePrimal(), x)  .- [42.881, 25.706]) < 0.01
end

@testset "Trivial QP" begin
    opt = PIPG.Optimizer()
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )

	P_ex = spdiagm([1/MathConstants.e, 2.0])
	q = SVector(1/MathConstants.e,2.0)
	H_ex = sparse([-1.0 -1.0])
	g = SVector(0.0)

    x = MOI.add_variables(model, 2)
    MOI.set(model,
    	MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}(),
    	MOI.ScalarQuadraticFunction(
    		MOI.ScalarQuadraticTerm{Float64}.([1/MathConstants.e, 2.0], x, x),
    		MOI.ScalarAffineTerm.([1/MathConstants.e,2.0], x), 0.0))
    MOI.add_constraint(model,
    	MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 1.0], x), 0.0),
    	MOI.GreaterThan(0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
	@test norm(MOI.get.(model, MOI.VariablePrimal(), x) .- [0.6893, -0.6893]) < 0.001
end

@testset "Basic QP" begin
    opt = PIPG.Optimizer()
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )

    x = MOI.add_variables(model, 3)
    MOI.set(model,
    	MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}(),
    	MOI.ScalarQuadraticFunction(
    		MOI.ScalarQuadraticTerm{Float64}.([5.0, -2, -1, 4, 3, 5],
    			[repeat([x[1]], 3); repeat([x[2]], 2); repeat([x[3]], 1)],
    			[x; x[2:3]; x[3]]),
    		MOI.ScalarAffineTerm.([2.0, -35, -47], x), 5.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
	@test norm(MOI.get.(model, MOI.VariablePrimal(), x) .- [3, 5, 7]) < 0.001
end

@testset "LP scalar 2" begin
    opt = PIPG.Optimizer(niters=100000000)
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    c = [50.0, 120.0]
    x = MOI.add_variables(model, length(c))
    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}(),
        MOI.ScalarQuadraticFunction(MOI.ScalarQuadraticTerm{Float64}[], MOI.ScalarAffineTerm.(c, x), 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([100.0, 200.0], x), 0.0),
        MOI.LessThan(10000.0))
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([10.0, 30.0], x), 0.0),
        MOI.LessThan(1200.0))
    sum_cstr = MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 1.0], x), 0.0),
        MOI.LessThan(60.0))
    lb_cstr = MOI.add_constraint(model,
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(0.3, x[1])], 0.0), MOI.GreaterThan(10.0))
    MOI.add_constraint(model, MOI.VectorOfVariables(x), MOI.Nonnegatives(2))

    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test norm(MOI.get.(model, MOI.VariablePrimal(), x)  .- [33.333, 26.666]) < 0.01
end
#=
prob = PIPG.Problem(
           PIPG.PTCone{Float64}((PIPG.POCone{Float64, 1}(), PIPG.NOCone{Float64, 5}())),
           PIPG.Reals{Float64, 2}(),
sparse([   0.3    0
         100.0  200.0
          10.0   30.0
           1.0    1.0
          -1.0    0
           0     -1.0]), spzeros(2,2) , MVector(-50.0, -120.0), MVector(10.0, 10000.0, 1200.0, 60.0, -0.0, -0.0), 0.0)
state = PIPG.State(prob, PIPG.LogDiagnostics{Float64}())
α = PIPG.compute_α(prob, 0.9)
PIPG.pipg(prob, state, 100000000, α, 1e-5, SVector{2, Float64}(zeros(2)), SVector{6, Float64}(zeros(6)))
# baseline run (opt) 15270124

prob = PIPG.Problem(
           PIPG.PTCone{Float64}((PIPG.POCone{Float64, 1}(), PIPG.NOCone{Float64, 3}())),
           PIPG.POCone{Float64, 2}(),
sparse([   0.3    0
         100.0  200.0
          10.0   30.0
           1.0    1.0]), spzeros(2,2) , MVector(-50.0, -120.0), MVector(10.0, 10000.0, 1200.0, 60.0), 0.0)
PIPG.State(prob, PIPG.LogDiagnostics{Float64}(), PIPG.GeoMean{Float64, 4, 2}())
α = PIPG.compute_α(prob, 0.9)
PIPG.pipg(prob, state, 100000000, α, 1e-5, SVector{2, Float64}(zeros(2)), SVector{4, Float64}(zeros(4)))

prob = PIPG.Problem(
           PIPG.PTCone{Float64}((PIPG.POCone{Float64, 1}(), PIPG.NOCone{Float64, 3}())),
           PIPG.POCone{Float64, 2}(),
sparse([ 0.3 0
         .1  .2
         .1   .3
         1.0    1.0]), spzeros(2,2) , MVector(-50.0, -120.0), MVector(10.0, 10.0, 12.0, 60.0), 0.0)
state = PIPG.State(prob, PIPG.LogDiagnostics{Float64}())
α = PIPG.compute_α(prob, 0.9)
PIPG.pipg(prob, state, 100000000, α, 1e-5, SVector{2, Float64}(zeros(2)), SVector{4, Float64}(zeros(4)))
=#

@testset "LP scalar modification" begin
    opt = PIPG.Optimizer()
    opt.niters *= 10
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    c = [50.0, 120.0]
    x = MOI.add_variables(model, length(c))
    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}(),
        MOI.ScalarQuadraticFunction(MOI.ScalarQuadraticTerm{Float64}[], MOI.ScalarAffineTerm.(c, x), 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([100.0, 200.0], x), 0.0),
        MOI.LessThan(10000.0))
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([10.0, 30.0], x), 0.0),
        MOI.LessThan(1200.0))
    sum_cstr = MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 1.0], x), 0.0),
        MOI.LessThan(110.0))
    lb_cstr = MOI.add_constraint(model,
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x[1])], 0.0), MOI.GreaterThan(0.0))
    MOI.add_constraint(model, MOI.VectorOfVariables(x), MOI.Nonnegatives(2))
    MOI.optimize!(model)
    MOI.modify(model, sum_cstr, MOI.ScalarConstantChange(60.0))
    MOI.modify(model, lb_cstr, MOI.ScalarConstantChange(10.0))
    MOI.optimize!(model)
    change = MOI.ScalarCoefficientChange(x[1], 0.3)
    MOI.modify(model, lb_cstr, change)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test norm(MOI.get.(model, MOI.VariablePrimal(), x)  .- [33.333, 26.666]) < 0.01
end

@testset "LP vector modification" begin
    opt = PIPG.Optimizer()
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    c = [50.0, 120.0]
    x = MOI.add_variables(model, length(c))
    MOI.set(model,
    	MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}(),
    	MOI.ScalarQuadraticFunction(MOI.ScalarQuadraticTerm{Float64}[], MOI.ScalarAffineTerm.(c, x), 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    cv = MOI.add_constraint(model,
    	MOI.VectorAffineFunction(MOI.VectorAffineTerm.(
    	[1,1,2,2,3,3],
    	[MOI.ScalarAffineTerm.([100.0, 200.0], x);
    	 MOI.ScalarAffineTerm.([10.0, 30.0], x);
    	 MOI.ScalarAffineTerm.([1.0, 1.0], x)]), [-10000.0, -1200.0, -110.0]),
    	MOI.Nonpositives(3))
    MOI.add_constraint(model, MOI.VectorOfVariables(x), MOI.Nonnegatives(2))
    MOI.optimize!(model)
    cst = MOI.VectorConstantChange([-10000.0, -1200.0, -60.0])
    MOI.modify(model, cv, cst)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test norm(MOI.get.(model, MOI.VariablePrimal(), x)  .- [30, 30]) < 0.01
    cst = MOI.VectorConstantChange([-10000.0, -600.0, -60.0])
    MOI.modify(model, cv, cst)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test norm(MOI.get.(model, MOI.VariablePrimal(), x)  .- [60, 0]) < 0.01
end

@testset "Geometric mean cone clone" begin
    opt = PIPG.Optimizer()
    #opt.scaling = PIPG.NoScaling
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    n = 3
    t = MOI.add_variable(model)
    x = MOI.add_variables(model, n)
    vov = MOI.VectorOfVariables([t; x])
    gmc = MOI.add_constraint(model, vov, MOI.GeometricMeanCone(n + 1))
    c = MOI.add_constraint(
        model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(Float64(1), x), Float64(0)),
        MOI.LessThan(Float64(n)),
    )
    MOI.set(
        model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(Float64(1), t)], Float64(0)),
    )
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.ObjectiveValue()) ≈ 1
end

@testset "Geometric mean cone clone2" begin
    opt = PIPG.Optimizer()
    opt.scaling = x->(PIPG.Equilibration(x), )
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    T = Float64
    n = 3
    n = 9
    t = MOI.add_variable(model)
    x = MOI.add_variables(model, n)
    vov = MOI.VectorOfVariables([t; x])
    gmc = MOI.add_constraint(model, vov, MOI.GeometricMeanCone(n + 1))
    cx =
        Vector{MOI.ConstraintIndex{MOI.ScalarAffineFunction{T},MOI.EqualTo{T}}}(
            undef,
            n,
        )
    for i in 1:n
        cx[i] = MOI.add_constraint(
            model,
            MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(T(1), x[i])], T(0)),
            MOI.EqualTo(T(1)),
        )
    end
    MOI.set(
        model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(T(1), t)], T(0)),
    )
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    MOI.optimize!(model)

    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
end

@testset "Infeasible SOC2" begin
    opt = PIPG.Optimizer()
    opt.scaling = x->(PIPG.Equilibration(x), )
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    T = Float64
    t = MOI.add_variable(model)
    x = MOI.add_variables(model, 2)
    MOI.add_constraints(model, x, MOI.GreaterThan.(T[3, 4]))
    MOI.set(model, MOI.ObjectiveFunction{MOI.VariableIndex}(), t)
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.DUAL_INFEASIBLE
end
@testset "Infeasible SOC" begin
    opt = PIPG.Optimizer()
    opt.scaling = x->(PIPG.GeoMean(x), )
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    T = Float64
    t = MOI.add_variable(model)
    x = MOI.add_variables(model, 2)
    MOI.add_constraints(model, x, MOI.GreaterThan.(T[3, 4]))
    c_soc = MOI.add_constraint(
        model,
        MOI.VectorOfVariables([t; x]),
        MOI.SecondOrderCone(3),
    )
    MOI.set(model, MOI.ObjectiveFunction{MOI.VariableIndex}(), t)
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    MOI.optimize!(model)
    @test isapprox(MOI.get(model, MOI.VariablePrimal(), t), T(5))
    MOI.delete(model, c_soc)
    MOI.optimize!(model)
    println(opt.problem)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.DUAL_INFEASIBLE
end

@testset "Quadratic test copies" begin
    opt = PIPG.Optimizer()
    opt.scaling = x->()
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    T = Float64
    obj_attr = MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{T}}()
    x = MOI.add_variables(model, 2)
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    vc1 = MOI.add_constraint(model, x[1], MOI.GreaterThan(T(1)))
    @test vc1.value == x[1].value
    vc2 = MOI.add_constraint(model, x[2], MOI.GreaterThan(T(2)))
    @test vc2.value == x[2].value
    # Basic model
    # min x^2 + y^2 | x>=1, y>=2
    MOI.set(
        model,
        obj_attr,
        MOI.ScalarQuadraticFunction(
            MOI.ScalarQuadraticTerm.(T(2), x, x),  # quad
            MOI.ScalarAffineTerm{T}[],  # affine terms
            T(0),  # constant
        ),
    )
    MOI.optimize!(model)
    println(opt.problem.P)
    variable_primal = [(x[1], T(1)), (x[2], T(2))]
    for (index, solution_value) in variable_primal
        @test isapprox(
            MOI.get(model, MOI.VariablePrimal(), index),
            solution_value
        )
    end
    @test isapprox(
            MOI.get(model, MOI.ObjectiveValue()),
            T(5),
        )
    MOI.set(
        model,
        obj_attr,
        MOI.ScalarQuadraticFunction(
            MOI.ScalarQuadraticTerm.(
                T[2, 1//4, 1//4, 1//2, 2],
                [x[1], x[1], x[2], x[1], x[2]],
                [x[1], x[2], x[1], x[2], x[2]],
            ),  # quad
            MOI.ScalarAffineTerm{T}[],  # affine terms
            T(0),  # constant
        ),
    )
    MOI.optimize!(model)
    println(opt.problem.P)
    variable_primal = [(x[1], T(1)), (x[2], T(2))]
    for (index, solution_value) in variable_primal
        @test isapprox(
            MOI.get(model, MOI.VariablePrimal(), index),
            solution_value
        )
    end
    @test isapprox(
            MOI.get(model, MOI.ObjectiveValue()),
            T(7),
        )
end

    const OPTIMIZER = MOI.instantiate(
        MOI.OptimizerWithAttributes(PIPG.Optimizer, MOI.Silent() => true),
    )

    const BRIDGED = MOI.instantiate(
        MOI.OptimizerWithAttributes(PIPG.Optimizer, MOI.Silent() => true),
        with_bridge_type = Float64,
    )
    const CONFIG = MOI.Test.Config(
        atol=1e-5,
        rtol=1e-5,
        optimal_status=MOI.OPTIMAL,
        exclude=Any[MOI.VariableName, MOI.delete,
            # not an integer solver
            MOI.ObjectiveBound, MOI.RelativeGap,
            # not simplex
            MOI.ConstraintBasisStatus, MOI.SimplexIterations,
            # not barrier, nor branch-and-cut
            MOI.BarrierIterations, MOI.NodeCount])

    MOI.Test.runtests(
        BRIDGED,
        CONFIG,
        exclude = [
            "test_attribute_NumberOfThreads",
            "test_linear_Indicator_",
            "test_model_UpperBoundAlreadySet"
        ],
        # This argument is useful to prevent tests from failing on future
        # releases of MOI that add new tests. Don't let this number get too far
        # behind the current MOI release though! You should periodically check
        # for new tests in order to fix bugs and implement new features.
        exclude_tests_after = v"0.10.5",
    )


text:
using Plasmo
using JuMP
using PlasmoAlgorithms

include("fisher.jl")

heur(g) = 16

r = lagrangesolve(g,update_method=:subgradient,max_iterations=30,lagrangeheuristic=heur)


text: println("hola"


text:
# # Ground
#
# LongwaveModePropagator treats the ground boundary of the Earth-ionosphere waveguide
# as a homogeneous, nonmagnetic material described by its electrical conductivity ``\sigma``
# and relative permittivity ``\epsilon_r``.
#
# In this example we'll compare the effect of ground conductivity on amplitude along
# the ground in the waveguide for day and nighttime ionospheres.

using Plots, Printf
using LongwaveModePropagator
const LMP = LongwaveModePropagator
nothing  #hide

# LWPC uses standard ground indices to describe combinations of relative permittivity
# and conductivity.
# The standard indices can be accessed from LMP as [`GROUND`](@ref).

sort(GROUND)

# We'll define all of the constant types and the day and night profiles.

## Constant values
const BFIELD = BField(50e-6, π/2, 0)

const TX = Transmitter(20e3)
const RX = GroundSampler(0:5e3:3000e3, Fields.Ez)

const DAY = Species(LMP.QE, LMP.ME, z->waitprofile(z, 75, 0.3), electroncollisionfrequency)
const NIGHT = Species(LMP.QE, LMP.ME, z->waitprofile(z, 82, 0.6), electroncollisionfrequency)
nothing  #hide

# We'll define a function to which we can pass the day or night `Species` and return the
# electric field amplitude.

function varyground(prf)
    amps = Vector{Vector{Float64}}(undef, length(GROUND))
    for i = 1:length(GROUND)
        wvg = HomogeneousWaveguide(BFIELD, prf, GROUND[i])
        E, a, p = propagate(wvg, TX, RX)
        amps[i] = a
    end
    return amps
end
nothing  #hide

# And here's a function to plot each of the curves.

function buildplots!(p, amps)
    cmap = palette(:thermal, length(GROUND)+1)  # +1 so we don't get to yellow

    for i = 1:length(GROUND)
        plot!(p, RX.distance/1000, amps[i];
              label=@sprintf("%d, %.1g", GROUND[i].ϵᵣ, GROUND[i].σ), color=cmap[i]);
    end
end
nothing  #hide

# First, the daytime results.

amps = varyground(DAY)

p = plot()
buildplots!(p, amps)
plot!(p; size=(600,400), ylims=(0, 95), title="Day", legend=(0.85, 1.02),
      xlabel="Range (km)", ylabel="Amplitude (dB)", legendtitle="ϵᵣ, σ")
#md savefig(p, "ground_day.png"); nothing # hide
#md # ![](ground_day.png)

# And now nighttime.

amps = varyground(NIGHT)

p = plot()
buildplots!(p, amps)
plot!(p; size=(600,400), ylims=(0, 95), title="Night", legend=(0.85, 1.02),
      xlabel="Range (km)", ylabel="Amplitude (dB)", legendtitle="ϵᵣ, σ")
#md savefig(p, "ground_night.png"); nothing # hide
#md # ![](ground_night.png)

# Low ground conductivity can have a significant influence on the signal propagation -
# there is strong attenuation.
# These low conductivities can be found in areas of sea or polar ice and industrial or
# city areas.
#
# The influence of ground conductivity on the received signal has similar impact on
# the day and night scenarios.


text:
# Clustering example following the tutorial:
# http://www.pointclouds.org/documentation/tutorials/cluster_extraction.php
# essentially same but implementation details are different.
# We use downsampled and plane removed cloud here for simplicity

using PCLCommon,PCLVisualization,PCLIO,PCLFilters,PCLSegmentation,Cxx

T = PointXYZ

pcd_file =joinpath(dirname(@__FILE__), "pcd", "2016-12-15-23-45-38-884_concat.pcd")
cloud = PointCloud{T}(pcd_file)

# Remove noise
sor = StatisticalOutlierRemoval{T}()
setMeanK(sor, 50)
setStddevMulThresh(sor, 1.0)
setInputCloud(sor, cloud)
filter(sor, cloud)

tree = PCLSearch.KdTree{T}()
setInputCloud(tree, cloud)

cluster_indices = icxx"std::vector<pcl::PointIndices>();"

ec = EuclideanClusterExtraction{T}()
PCLSegmentation.setClusterTolerance(ec, 0.02)
PCLSegmentation.setMinClusterSize(ec, 100)
PCLSegmentation.setMaxClusterSize(ec, 25000)
PCLSegmentation.setSearchMethod(ec, tree)
setInputCloud(ec, cloud)
@time extract(ec, cluster_indices)

clusters = []
for i in 0:Int(length(cluster_indices))-1
    c = PointCloud{T}()
    icxx"""
    auto cloudp = $(c.handle);
    for (auto idx : $(cluster_indices)[$i].indices) {
        cloudp->push_back($(cloud.handle)->points[idx]);
    }
    cloudp->width = cloudp->points.size();
    cloudp->height = 1;
    cloudp->is_dense = true;
    """
    push!(clusters, c)
end

@show length(clusters)
@assert length(clusters) > 0

if isdefined(:vis) && vis
    viewer = PCLVisualizer("pcl visualizer")
    addCoordinateSystem(viewer, 1.2, 0,0,0)
    addPointCloud(viewer, cloud, PointCloudColorHandlerCustom(cloud, 255, 0, 0), id="cloud")
    for (idx, cluster) in enumerate(clusters)
        addPointCloud(viewer, cluster,
            PointCloudColorHandlerCustom(cluster, rand(0:255),rand(0:255),rand(0:255)),
            id="cluster #$idx")
    end
    spin(viewer)
    close(viewer); viewer = 0; gc();gc()
end


text:
function simple1(a::Int, b::Int)::Int
    c::Int = a + b

    if c > 5
        c = 2 * c
    else
        c = 3 * c
    end

    return c
end

function multiply(a::Int, b::Int)::Int
    return a * b
end

function divide(a::Int, b::Int)::Int
    return a / b
end

function simple2(a::Int)::Int
    return multiply(a, 2)
end

function simple3(a::Int)::Int
    return divide(a, 2)
end

function simple4(a::Int)::Int
    return multiply(a, 3)
end

function simple5(limit::Int)::Int
    sum::Int = 0

    for i = 1:limit
        sum += i
    end

    return sum
end

function array_ops(size::Int)::Array{Int, 1}
    a::Array{Int, 1} = ones(size)

    b::Array{Int, 1} = a .* 2

    b = b .+ 2

    b = b .- 2

    b = b ./ 2

    return b
end

max_stars_repo_name: phrb/PDP

max_issues_repo_name: phrb/PDP

max_forks_repo_name: phrb/PDP


text:
println("Hello, World!")

# This is a comment

#= This is a multi-line comment
The second line of the comment =#

# println ("this is not allowed")
println("This is allowed.")
println("This is also allowed.");

println(1 + 2)
1 + 2 # Will not print in script

# Installing, updating and removing packages
using Pkg # Alternatively, enter pkg mode in REPL with ]
Pkg.add("Example")
Pkg.status()
Pkg.update("Example") # Pkg.update() updates all packages
Pkg.rm("Example") # removes from environment, does not delete files

# Environments
# In pkg mode:
# activate tutorial
# st
# add Example
# activate # goes back to default environment


text:
using Zygote

## Simple gradients

gradient(x->x^2, 3)

## Gradients of ... what?!

fs = Dict("sin" => sin, "cos" => cos, "tan" => tan);

gradient(x -> fs[readline()](x), 1)


text:
function lgd100()
#
#  test polynomial suggested by Goedecker
#   Legendre polynomial of degree 100
#
    p = lgd(100);
    z = [-.9997137267734412; -.9984919506395958; -.9962951347331251; -.9931249370374435;
         -.9889843952429917; -.9838775407060570; -.9778093584869183; -.9707857757637063;
         -.9628136542558155; -.9539007829254917; -.9440558701362560; -.9332885350430795;
         -.9216092981453340; -.9090295709825297; -.8955616449707270; -.8812186793850184;
         -.8660146884971646; -.8499645278795913; -.8330838798884008; -.8153892383391763;
         -.7968978923903145; -.7776279096494955; -.7575981185197072; -.7368280898020207;
         -.7153381175730564; -.6931491993558020; -.6702830156031410; -.6467619085141293;
         -.6226088602037078; -.5978474702471787; -.5725019326213812; -.5465970120650942;
         -.5201580198817631; -.4932107892081909; -.4657816497733580; -.4378974021720315;
         -.4095852916783015; -.3808729816246300; -.3517885263724217; -.3223603439005292;
         -.2926171880384720; -.2625881203715035; -.2323024818449740; -.2017898640957360;
         -.1710800805386033; -.1402031372361140; -.1091892035800611; -.7806858281343664e-1;
         -.4687168242159163e-1; -.1562898442154308e-1; .1562898442154308e-1; .4687168242159163e-1;
         .7806858281343664e-1; .1091892035800611; .1402031372361140; .1710800805386033;
         .2017898640957360; .2323024818449740; .2625881203715035; .2926171880384720;
         .3223603439005292; .3517885263724217; .3808729816246300; .4095852916783015;
         .4378974021720315; .4657816497733580; .4932107892081909; .5201580198817631;
         .5465970120650942; .5725019326213812; .5978474702471787; .6226088602037078;
         .6467619085141293; .6702830156031410; .6931491993558020; .7153381175730564;
         .7368280898020207; .7575981185197072; .7776279096494955; .7968978923903145;
         .8153892383391763; .8330838798884008; .8499645278795913; .8660146884971646;
         .8812186793850184; .8955616449707270; .9090295709825297; .9216092981453340;
         .9332885350430795; .9440558701362560; .9539007829254917; .9628136542558155;
         .9707857757637063; .9778093584869183; .9838775407060570; .9889843952429917;
         .9931249370374435; .9962951347331251; .9984919506395958; .9997137267734412]
    z = [z ones(100)]
    # Legendre polynomial, roots are sensitive, see plot.

    p, PolyZeros(z)
end


text:
#if lowercase(get(ENV, "CI", "false")) == "true"
#    include("install_dependencies.jl")
#end

#using CausalityToolsBase
using Test, TransferEntropy, Entropies
KDTree = Entropies.KDTree
BruteForce = Entropies.BruteForce

@testset "Mutual information" begin
    x = rand(100)
    y = rand(100)
    z = Dataset(rand(100, 2))
    w = Dataset(rand(100, 3))

    # Estimators for which Renyi entropies can be used
    est = VisitationFrequency(RectangularBinning(0.2))
    @test mutualinfo(x, y, est) isa Real
    @test mutualinfo(x, y, est, q = 2) isa Real
    @test mutualinfo(x, z, est) isa Real
    @test mutualinfo(z, x, est) isa Real
    @test mutualinfo(z, w, est) isa Real

    # Estimators for which Renyi entropies cannot be used
    est_kl = KozachenkoLeonenko()
    est_k = Kraskov()
    est_k1 = Kraskov1(2)
    est_k2 = Kraskov2(2)

    @test mutualinfo(x, y, est_kl) isa Real
    @test mutualinfo(x, y, est_kl, base = 2) isa Real
    @test mutualinfo(x, z, est_kl) isa Real
    @test mutualinfo(z, x, est_kl) isa Real
    @test mutualinfo(z, w, est_kl) isa Real

    @test mutualinfo(x, y, est_k) isa Real
    @test mutualinfo(x, y, est_k, base = 2) isa Real
    @test mutualinfo(x, z, est_k) isa Real
    @test mutualinfo(z, x, est_k) isa Real
    @test mutualinfo(z, w, est_k) isa Real

    @test mutualinfo(x, y, est_k1) isa Real
    @test mutualinfo(x, y, est_k1, base = 2) isa Real
    @test mutualinfo(x, z, est_k1) isa Real
    @test mutualinfo(z, x, est_k1) isa Real
    @test mutualinfo(z, w, est_k1) isa Real

    @test mutualinfo(x, y, est_k2) isa Real
    @test mutualinfo(x, y, est_k2, base = 2) isa Real
    @test mutualinfo(x, z, est_k2) isa Real
    @test mutualinfo(z, x, est_k2) isa Real
    @test mutualinfo(z, w, est_k2) isa Real
end

@testset "Transfer entropy" begin
    s, t, c = rand(100), rand(100), rand(100)

    println("Starting transfer entropy tests...")
    @testset "Generalized Renyi transfer entropy" begin
        # Straight-forward estimators
        est_vf = VisitationFrequency(RectangularBinning(4))
        ests = [
            est_vf,
            TransferOperator(RectangularBinning(5)),
            Hilbert(source = Phase(), target = Amplitude(), est_vf),
            Hilbert(source = Amplitude(), target = Amplitude(), est_vf),
            Hilbert(source = Phase(), target = Phase(), est_vf),
            Hilbert(source = Phase(), target = Phase(), cond = Amplitude(), est_vf),
            NaiveKernel(0.5, KDTree),
            NaiveKernel(0.5, BruteForce)
        ]

        @testset "Generalized Renyi transfer entropy $(ests[i])"  for i in 1:length(ests)
            println(ests[i])
            est = ests[i]
            te =  transferentropy(s, t, est, q = 2, base = 2)
            tec = transferentropy(s, t, c, est, q = 1, base = 3)

            @test te isa Real
            @test tec isa Real

            if est isa VisitationFrequency
                @test te > 0 || te ≈ 0
                @test tec > 0 || tec ≈ 0
            end
        end

        @testset "SymbolicPermutation" begin
            m, τ = 3, 1
            est = SymbolicPermutation(m = m, τ = τ)
            te =  transferentropy(s, t, est, q = 2, base = 2)
            tec = transferentropy(s, t, c, est, q = 1, base = 3)

            S, T, C = ([zeros(Int, length(s)-(m-1)*τ) for x in 1:3]...,)
            te_inplace =  transferentropy!(S, T, s, t, est, q = 2, base = 2)
            tec_inplace = transferentropy!(S, T, C, s, t, c, est, q = 1, base = 3)

            @test te isa Real
            @test tec isa Real
            @test te_inplace isa Real
            @test tec_inplace isa Real
        end
    end

    @testset "Shannon transfer entropy" begin
        ests = [
            Kraskov(),
            KozachenkoLeonenko()
        ]
        @testset "Nearest neighbors $(ests[i])"  for i in 1:length(ests)
            println(ests[i])
            est = ests[i]
            te =  transferentropy(s, t, est, base = 2)
            tec = transferentropy(s, t, c, est, base = 3)
            @test te isa Real
            @test tec isa Real

        end
    end

    @testset "Automated estimators" begin
        @testset "Variable exclusion" begin
             # Use periodic signals, so we also can test variable selection methods,
            # which for sensible testing, need to have their autocorrelation minima
            # > 1.
            s = sin.(1:100) .+ rand(100)
            t = sin.(1:100) .+ rand(100)
            c = sin.(1:100) .+ rand(100)

            τexclude = 1
            vars = TransferEntropy.construct_candidate_variables([s], [t], τexclude = nothing)
            vars_ex = TransferEntropy.construct_candidate_variables([s], [t], τexclude = τexclude)
            fvars = Iterators.flatten(vars[1][1:end-1]) |> collect
            fvars_ex = Iterators.flatten(vars_ex[1][1:end-1]) |> collect
            @test τexclude ∈ abs.(fvars)
            @test τexclude ∉ abs.(fvars_ex)
            @test length(fvars) > length(fvars_ex)

            vars = TransferEntropy.construct_candidate_variables([s], [t], [c], τexclude = nothing)
            vars_ex = TransferEntropy.construct_candidate_variables([s], [t], [c], τexclude = τexclude)
            fvars = Iterators.flatten(vars[1][1:end-1]) |> collect
            fvars_ex = Iterators.flatten(vars_ex[1][1:end-1]) |> collect
            @test τexclude ∈ abs.(fvars)
            @test τexclude ∉ abs.(fvars_ex)
            @test length(fvars) > length(fvars_ex)
        end

        est = VisitationFrequency(RectangularBinning(3))
        te_st, params_st = bbnue(s, t, est)
        te_stc, params_stc = bbnue(s, t, c, est)
        @test te_stc isa Float64
        @test te_stc >= 0.0
    end

end


text:
using Pkg
Pkg.instantiate()

using PackageCompiler

# create_sysimage([:CSV,:DataFrames,:FileIO,:MeshIO,:GaussianProcesses,:GeometryBasics,:MVApp];precompile_execution_file="precompile_app.jl",sysimage_path="MVApp_plus.so")
create_sysimage([:CSV,:DataFrames,:FileIO,:MeshIO,:GaussianProcesses,:GeometryBasics,:MVApp];
                precompile_execution_file="precompile_app.jl",
                replace_default=true)


text:
using SafeTestsets, Test

files = [
    "multivariate_normal_tests.jl",
    "blocking_tests.jl",
    "optimization_tests.jl",
    "utility_tests.jl"
]

for f in files
    include(f)
end

@safetestset "gaussian tests" begin
    include("gaussian_tests.jl")
end

@safetestset "lognormal race tests" begin
    include("lognormal_race_tests.jl")
end


text:
module AppendedArraysTests

using Test
using FillArrays
using Gridap.Arrays

a = collect(Float64,11:20)
b = collect(Float64,101:120)
c = lazy_append(a,b)
r = vcat(a,b)
test_array(c,r)

@test sum(c) == sum(r)
@test ∑(c) == sum(r)

d = lazy_map(-,c)
r = -c
test_array(d,r)
@test isa(d,AppendedArray)

e = lazy_map(-,c,d)
r = c-d
test_array(e,r)
@test isa(e,AppendedArray)

d = lazy_map(-,Float64,c)
r = -c
test_array(d,r)
@test isa(d,AppendedArray)

e = lazy_map(-,Float64,c,d)
r = c-d
test_array(e,r)
@test isa(e,AppendedArray)

a = collect(11:20)
b = collect(Float64,101:120)
c = lazy_append(a,b)
d = lazy_map(i->Float64(i),c)
r = vcat(Float64.(a),b)
test_array(d,r)
#print_op_tree(d)

a1 = [1,2,3]
b1 = [4,5]
c1 = lazy_append(a1,b1)

a2 = [1,2]
b2 = [3,4,5]
c2 = lazy_append(a2,b2)

d = lazy_map(+,c1,c2)
test_array(d,c1+c2)
#print_op_tree(d)

# test when lazy_array leads to an array with an eltype that is not concrete.

struct FooFun{T}
  f::T
end
Arrays.evaluate!(cache,f::FooFun,args...) = f.f(args...)

a1 = Fill(FooFun(+),3)
b1 = Fill(FooFun(-),2)
c1 = lazy_append(a1,b1)
c2a = CompressedArray([3.0],fill(1,3))
c2b = CompressedArray([3.0],fill(1,2))
c2 = lazy_append(c2a,c2b)
d = lazy_map(evaluate,c1,c2)
test_array(d,[3,3,3,-3,-3])

a1 = Fill(FooFun(+),3)
b1 = Fill(FooFun(-),2)
c1 = lazy_append(a1,b1)
c2 = Fill(3.0,5)
d = lazy_map(evaluate,c1,c2)
test_array(d,[3,3,3,-3,-3])

a1 = Fill(FooFun(+),3)
b1 = Fill(FooFun(-),2)
c1 = lazy_append(a1,b1)
c2 = CompressedArray([3.0],fill(1,5))
d = lazy_map(evaluate,c1,c2)
test_array(d,[3,3,3,-3,-3])

end # module


text:
#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct Deterministic <: Forecast
        label::String
        data::TimeSeries.TimeArray
    end

A deterministic forecast for a particular data field in a Component.

# Arguments
- `label::String`: label of component parameter forecasted
- `data::TimeSeries.TimeArray`: timestamp - scalingfactor
"""
mutable struct Deterministic <: Forecast
    "label of component parameter forecasted"
    label::String
    "timestamp - scalingfactor"
    data::TimeSeries.TimeArray
end

function Deterministic(; label, data, )
    Deterministic(label, data, )
end

"""Get Deterministic label."""
get_label(value::Deterministic) = value.label
"""Get Deterministic data."""
get_data(value::Deterministic) = value.data

"""Set Deterministic label."""
set_label!(value::Deterministic, val::String) = value.label = val
"""Set Deterministic data."""
set_data!(value::Deterministic, val::TimeSeries.TimeArray) = value.data = val


text:
# Make SQL fun!

module FunSQL

using Dates
using PrettyPrinting: PrettyPrinting, pprint, quoteof, tile_expr, literal
using OrderedCollections: OrderedDict, OrderedSet

const SQLLiteralType =
    Union{Missing, Bool, Number, AbstractString, Dates.AbstractTime}

"""
    render(node::Union{SQLNode, SQLClause}; dialect = :default)::SQLStatement

Convert the given SQL node or clause object to a SQL string.
"""
function render
end

"""
Base error class for all errors raised by FunSQL.
"""
abstract type FunSQLError <: Exception
end

include("dissect.jl")
include("dialects.jl")
include("types.jl")
include("statements.jl")
include("entities.jl")
include("quote.jl")
include("clauses.jl")
include("nodes.jl")
include("annotate.jl")
include("translate.jl")
include("render.jl")

const Not = Fun.not
const And = Fun.and
const Or = Fun.or
const Like = Fun.like
const In = Fun."in"
const NotIn = Fun."not in"
const IsNull = Fun."is null"
const IsNotNull = Fun."is not null"
const Case = Fun.case
const Concat = Fun.concat
const Coalesce = Fun.coalesce

end


text:
using LocallyCompetitive
using Test
using Random
using Statistics

# Testing params
M = 100
N = 100

k = 10

Random.seed!(0)

# Dictionary
basis = randn(M,N)

# Make k sparse signal
vals  = rand(M)
vals[sortperm(rand(M))[k+1:end]] .= 0

# Input signal with k sparse structure
sig = basis*vals

# Learn 'vals' based upon sig and the basis
out = LocallyCompetitive.lca(sig, basis)

recon = basis*out

sparse_orig = vals .!= 0
sparse_new  = out  .!= 0

sse_orig  = mean(recon.^2)
sse_new   = mean(sig.^2)
sse_diff  = mean((sig.-recon).^2)

# Same sparsity found in output
@test sparse_orig == sparse_new

# Reconstruction models data
@test sse_orig*0.01 > sse_diff
@test sse_new *0.01 > sse_diff


text:
macro src()
    @static if VERSION >= v"0.7-" && length(:(@test).args) == 2
        esc(quote
            (__module__,
             __source__.file == nothing ? "?" : String(__source__.file),
             __source__.line)
        end)
    else
        esc(quote
            (current_module(),
             (p = Base.source_path(); p == nothing ? "REPL" : p),
             Int(unsafe_load(cglobal(:jl_lineno, Cint))))
        end)
    end
end

macro log(stmt)
    # "[HTTP]: Connecting to remote host..."
    return esc(:(verbose && (write(logger, "[HTTP - $(rpad(Dates.now(), 23, ' '))]: $($stmt)\n"); flush(logger))))
end

macro catcherr(etype, expr)
    esc(quote
        try
            $expr
        catch e
            isa(e, $etype) ? e : rethrow(e)
        end
    end)
end


text:
@testset "util" begin
    @testset "eachrow" begin
        X = randn(2, 5)

        y, pb = Zygote.pullback(collect ∘ eachrow, X)
        @test y == collect(eachrow(X))

        # Just make the seed the output since it's a valid cotangent for itself.
        ȳ = y

        # Compute pullback using Zygote and FiniteDifferences and test they roughly agree.
        dX_fd = FiniteDifferences.j′vp(central_fdm(5, 1), collect ∘ eachrow, ȳ, X)
        dX_ad = pb(ȳ)

        @test first(dX_fd) ≈ first(dX_ad)
    end
    @testset "inv(::Diagonal)" begin
        x = randn(3)

        f = diag ∘ inv ∘ Diagonal

        y, pb = Zygote.pullback(f, x)
        @test y == f(x)

        ȳ = randn(3)

        dX_fd = FiniteDifferences.j′vp(central_fdm(5, 1), f, ȳ, x)
        dX_ad = pb(ȳ)

        @test first(dX_fd) ≈ first(dX_ad)
    end
    @testset "inv(::Cholesky{<:BlasFloat, <:StridedMatrix}" begin
        A = randn(3, 3)

        f = A -> inv(cholesky(A * A' + I))

        Ω, pb_inv_Chol = Zygote.pullback(f, A)
        @test Ω == f(A)

        ΔΩ = randn(3, 3)

        dX_fd = FiniteDifferences.j′vp(central_fdm(5, 1), f, ΔΩ, A)
        dX_ad = pb_inv_Chol(ΔΩ)

        @test first(dX_fd) ≈ first(dX_ad)
    end
end


text:
using Lorenz96
using SoftPosit
using PyPlot

Xp = L96(Posit16,F=8.0,N=1000,n=256,Δt=0.05)
Xf = L96(Float64,F=8.0,N=1000,n=256,Δt=0.05)

##
ioff()

cmin = -10
cmax = 10

fig,(ax1,ax2) = subplots(1,2,sharex=true,sharey=true,figsize=(8,3))

ax1.imshow(Xf,vmin=cmin,vmax=cmax,interpolation="bilinear",aspect="auto",origin="lower")
ax2.pcolormesh(Xp,vmin=cmin,vmax=cmax)

ax1.set_xticks([])
ax1.set_yticks([])

xlim(0,size(Xf)[2])
ylim(0,size(Xf)[1])

ax1.set_title("Float64",loc="left")
ax2.set_title("Posit16",loc="left")

ax1.set_ylabel(L"$X_i$")
ax1.set_xlabel("time")
ax2.set_xlabel("time")

tight_layout()
savefig("figs/hovmoeller.png",dpi=200)


text:
############################################################ Sentido 1: México-Acapulco #############################################################

function Sense1()
   Sense_1 = Highway2D(2, 3640+6)
    return Sense_1, Sense_1.highway
end

Sections1_CR = [150, 350, 535, 750, 1000,
                1250, 1490, 1600, 1770, 2000,
                2140, 2340, 2560, 2740, 2940,
                3100, 3300, 3640]

#(xramp, lramp)
S1_ramp = [(371, 38), (557, 3), (1015, 31), (1512, 100),
           (1628, 15), (1791, 10), (2158, 7), (2358, 20),
           (2580, 10), (2958, 6), (3115, 10)]

S1_ramp_local = [(371, 38), (557, 3), (1015, 31), (1512, 100),
          (1628, 15), (1791, 10), (2158, 7), (2358, 20),
          (2358, 20), (2580, 10), (2958, 6), (3115, 10)]

S1_ramp_ampli = [(371, 38), (2358, 20), (2580, 10),
                 (2958, 6), (3115, 10)]

#(p2, p1)
P_ramp_S1 = [
################  12 hours
                (0.445,	0.056),#0.056
                #(0.,	0.001),      # Entrada a Gasolinera
                #(0.378,	0.022),    # Entrada de Paloma
                (0.096,	0.009),    # Entrada de Paloma
                (0.096,	0.009),    # Entrada de Vicente Guerrero
                (0.095,	0.016),    # Entrada de Ahuatepec
                (0.038,	-0.01),    # Entrada de Diana
                (0.067,	0.017),    # Entrada entre Diana y Plan de Ayala
                (0.056,	0.007),    # Entrada después de Plan de Ayala
                (-0.053,	-0.007), # Salida antes de Tabachines
                (-0.144,	-0.019), # Salida de Morelos
                (-0.332,	-0.014), # Salida carretera libre
                (0.036,	-0.003),   # Entrada de Burgos
                #(0.01,	0.008),    # Calle entre Burgos y Brisas
                (-0.101,	-0.02)   # Salida de Brisas
                ]

P_local_S1 = [
################  12 hours
                (0.445,	0.056),
                #(0.378,	0.022),    # Entrada de Paloma
                (0.096,	0.009),    # Entrada de Paloma
                (0.096,	0.009),    # Entrada de Vicente Guerrero
                (0.095,	0.016),    # Entrada de Ahuatepec
                (0.038,	-0.01),    # Entrada de Diana
                (0.067,	0.017),    # Entrada entre Diana y Plan de Ayala
                (0.056,	0.007),    # Entrada después de Plan de Ayala
                (-0.053,	-0.007), # Salida antes de Tabachines
                (-0.144,	-0.019), # Salida a Morelos
                (-0.011, 0.006),   # Salida al puente
                (-0.332,	-0.014), # Salida carretera libre
                (0.036,	-0.003),   # Entrada de Burgos
                (-0.101,	-0.02)   # Salida de Brisas
                ]

P_ampli_S1 = [
#################  12 hours
              (0.445,	0.056),    # Flujo inicial
              (0.096,	0.009),    # Entrada de Paloma
              #(0.378,	0.022),    # Entrada de Paloma
                #(0.,	0.001)
                #(0.237,	0.018)
                #(0.01,	0.008)
              (0.155, 0.013),    # Antes del puente
              (-0.332,	-0.014), # Salida carretera libre (después del punte)
              (0.036,	-0.003),   # Entrada de Burgos
              (-0.101,	-0.02)   #Salida de Brisas
              ]

############################################################## Sentido 2: Acapulco-México ##########################################################

function Sense2()
    Sense_2 = Highway2D(3, 3640+6)
    for i = 1:3440 # 3685 = 3863 - 177 - 1
        Obstacle!(Sense_2.highway[1].highway[i])
    end
    return Sense_2, Sense_2.highway
end

Sections2_CR = [150, 350, 550, 790, 880,
                1100, 1300, 1480, 1660, 1900,
                2080, 2200, 2320, 2550, 2850,
                3050, 3290, 3520, 3640]

#(xramp, lramp)
S2_ramp = [(562, 35), (810, 5), (900, 4), (1500, 20),
           (1686, 5), (2107, 4), (2342, 7), (2871, 22),
           (3314, 4), (3542, 5)]

 S2_ramp_local = [(562, 35), (810, 5), (900, 4), (1500, 20),
            (1686, 5), (2107, 4), (2342, 7), (2871, 22),
            (3314, 4), #=(3450, 20),=# (3542, 5)]

S2_ramp_ampli =  [(562, 35), (810, 5), (900, 4), (1500, 20),
                     (3450, 20), (3542, 5)]

P_ramp_S2 = [
#################  12 hours
                (0.521,	0.08),     # Flujo Inicial
                (0.068,	-0.01),    # ITESM/Emiliano Zapata
                (0.073,	0.019),    # Brisas
                #(-0.026,	-0.013), #### Entre Brisas y Burgos
                (-0.02,	0.003),    # Burgos
                #(0.269,	0.02),   #### Entrada carretera libre
                (0.24,	0.033),    # Después del puente
                (0.076,	-0.003),   # Tabachines
                (-0.077,	-0.003), # Plan de Ayala
                #(-0.048,	-0.014), #### Entre Plan de Ayala y Diana
                (-0.053,	-0.009), #  Diana
                (-0.121,	-0.002), # Ahuatepec
                (-0.063,	-0.005), # Vicente Guerrero
                (-0.063,	-0.005)  # Paloma
                #(-0.417,	-0.021), # Paloma
                #(0.02,	0.004)     # Gasolinera
            ]

P_ampli_S2 = [
#################  12 hours
                (0.521,	0.08),     # Flujo Inicial
                (0.068,	-0.01),    # ITESM/Emiliano Zapata
                (0.073,	0.019),    # Brisas
                #(-0.026,	-0.013), #### Entre Brisas y Burgos
                (-0.02,	0.003),    # Burgos
                #(0.269,	0.02),   #### Entrada carretera libre
                (0.24,	0.033),    # Después del puente
                (-0.238, -.022),   # Termina Express
                (-0.063,	-0.005)  # Paloma
                #(-0.417,	-0.021), # Paloma
              ]

P_local_S2 = [
#################  12 hours
                (0.521,	0.08),     # Flujo Inicial
                (0.068,	-0.01),    # ITESM/Emiliano Zapata
                (0.073,	0.019),    # Brisas
                #(-0.026,	-0.013), #### Entre Brisas y Burgos
                (-0.02,	0.003),    # Burgos
                #(0.269,	0.02),   #### Entrada carretera libre
                (0.24,	0.033),    # Después del puente
                (0.076,	-0.003),   # Tabachines
                (-0.077,	-0.003), # Plan de Ayala
                #(-0.048,	-0.014), #### Entre Plan de Ayala y Diana
                (-0.053,	-0.009), #  Diana
                (-0.121,	-0.002), # Ahuatepec
                (-0.063,	-0.005), # Vicente Guerrero
                #(0.238, 0.22),      # Entrada Express
                (-0.063,	-0.005)  # Paloma
                #(-0.417,	-0.021) # Paloma
                #(0.02,	0.004)     # Gasolinera
            ]


text:
function seed(seed)
    Random.seed!(seed)
end

random_int(;min=0::Int, max=9999::Int)="$(rand(min : max))"
random_digit()= "$(rand(0:9))"
random_digit_not_null()= "$(rand(1:9))"

function random_digit_or_empty()
    if (rand(0:1) == 1)
        return "$(rand(0:9))"
    else
        return " "
    end
end

function random_digit_not_null_or_empty()
 if (rand(0:1)==1)
    return "$(rand(1:9))"
else
    return " "
end
end

random_number(;digits= 1::Int ) = "$(rand(0 : (10^digits) - 1))"
random_letter()= (rand(0:1) >= 0.5) ? "$(rand('a':'z'))" : "$(rand('A':'Z'))"
random_element(elements=("a", "b", "b"))= elements[rand(1: length(elements))]

function randomize_nb_elements(;number=10::Int, le=false::Bool, ge=false::Bool)
    if le==true && ge==true
        return "$(number)"
    end
    if ge==true
        min = 100
    else
        min=60
    end
    if le==true
        max = 100
    else
        max=140
    end
    "$(((rand(min : max))number / 100) + 1 )"
end

function numerify( text = "####"::String)
    for i in text
        text= replace(text,"#"=>"$(rand(0:9))",count=1)
    end
    return text
end

lexify( text="????"::String)= (for i=text  text= replace(text,"?"=>"$(random_letter())",count=1) end ; return text)
bothify( text="## ??"::String)=lexify(numerify(text))


text:
using DSGE
using HDF5
path = dirname(@__FILE__)
include("../util.jl")

# Test in model optimization
m = Model990()
m.testing=true

file = h5open("$path/../reference/optimize.h5","r")
x0 = read(file, "params")
data = read(file, "data")
minimum_ = read(file, "minimum")
f_minimum = read(file, "f_minimum")
H_expected = read(file, "H")
close(file)

# See src/estimate/estimate.jl
update!(m, x0)
n_iterations = 3

@time out, H = optimize!(m, data; iterations=n_iterations)

@test_matrix_approx_eq minimum_ out.minimum
@test_approx_eq_eps f_minimum out.f_minimum 5e-7
@test_matrix_approx_eq H_expected H

nothing


text:
function weird_model(t::Trace)
    choice = sample(t, "choice", Beta(1.0, 1.0))
    if choice > 0.5
        sample(t, "branch1", Normal())
    else
        sample(t, "branch2", LogNormal())
    end
end

@testset "table creation" begin
    sqm = SQLModel(weird_model; num_iterations=500)
    q = """
        SELECT
            AVG(branch1_logprob) as avg_1_logprob,
            AVG(branch2_logprob) as avg_2_logprob
        FROM
            weird_model_results
        WHERE
            choice < 0.25;
    """
    result = query(sqm, q)
    @info result
end

function regression(t::Trace, X::Vector{Float64}, y::Vector{Float64})
    alpha = sample(t, "alpha", Normal())
    beta = sample(t, "beta", Normal())
    sigma = sample(t, "sigma", LogNormal())
    mu = alpha .+ X .* beta
    observe(t, "obs", MvNormal(mu, sigma), y)
end

@testset "query posterior" begin
    N = 50
    X = randn(N,)
    true_alpha = 2.0
    true_beta = -0.3
    true_sigma = 0.25
    y = true_alpha .+ X .* true_beta .+ randn(N,) .* true_sigma
    train_N = 25
    sqm = SQLModel(
        regression;
        params = (X[1:train_N], y[1:train_N]),
        method = METROPOLIS,
        num_iterations = 21000,
        inference_params = Dict("burn" => 1000, "thin" => 100)
    )
    # examine posterior coefficients for "executive consumption"
    q = """
        SELECT
            AVG(alpha) as alpha_mean,
            AVG(beta) as beta_mean,
            AVG(sigma) as noise_level
        FROM regression_results;
    """
    @info query(sqm, q)
end


text:
const dexlang_root = dirname(dirname(@__DIR__))
cd(dexlang_root) do
    run(`make build-ffis`)
end

max_stars_repo_name: justinchiu/dex-lang

max_issues_repo_name: justinchiu/dex-lang

max_forks_repo_name: justinchiu/dex-lang


text:
abstract type AbstractGaussianModel <: AbstractSinglePotential end

### struct for defining a Gaussian force model
mutable struct GaussianModel <: AbstractGaussianModel
  dim::Int64
  nbasis::Int64
  shape::Tuple{Vararg{Int64}}
  mus::Vector{Vector{Float64}}
  sigs::Vector{Vector{Float64}}
  theta::Matrix{Float64}
  f::Vector{Float64}
  em::Vector{Float64}
  function GaussianModel(shape::Tuple{Vararg{Int64}},mus::Vector{Vector{Float64}},sigs::Vector{Vector{Float64}},
                         theta::Matrix{Float64})
    dim = length(shape)
    nbasis = reduce(*,shape,init=1)
    new(dim,nbasis,shape,mus,sigs,theta,zeros(Float64,dim),zeros(Float64,nbasis))
  end
end

function GaussianModel(shape::Tuple{Vararg{Int64}},mus::Vector{Vector{Float64}},sigs::Vector{Vector{Float64}})
  dim::Int64 = length(shape)
  nbasis::Int64 = reduce(*,shape,init=1)
  theta::Matrix{Float64} = zeros(Float64, (nbasis,dim))
  GaussianModel(shape,mus,sigs,theta)
end

"""
struct for defining a time-dependent Gaussian force model
make sure the last set of basis functions is the time-dependence
"""
mutable struct GaussianModelTD <: AbstractGaussianModel
  dim::Int64
  nbasis::Int64
  shape::Tuple{Vararg{Int64}}
  mus::Vector{Vector{Float64}}
  sigs::Vector{Vector{Float64}}
  theta::Matrix{Float64}
  f::Vector{Float64}
  em::Vector{Float64}
  function GaussianModelTD(shape::Tuple{Vararg{Int64}},mus::Vector{Vector{Float64}},sigs::Vector{Vector{Float64}},
                         theta::Matrix{Float64})
    dim = length(shape)-1
    nbasis = reduce(*,shape,init=1)
    new(dim,nbasis,shape,mus,sigs,theta,zeros(Float64,dim),zeros(Float64,nbasis))
  end
end

function GaussianModelTD(shape::Tuple{Vararg{Int64}},mus::Vector{Vector{Float64}},sigs::Vector{Vector{Float64}})
  dim::Int64 = length(shape)-1
  nbasis::Int64 = reduce(*,shape,init=1)
  theta::Matrix{Float64} = zeros(Float64, (nbasis,dim))
  GaussianModelTD(shape,mus,sigs,theta)
end

"""
Helper functions that get useful information about the model
"""
function getdimensionality(gm::AbstractGaussianModel)
  return gm.dim
end

# shape of the jacobian with respect to the parameters
function jacshape(gm::AbstractGaussianModel)
  return (gm.dim*gm.nbasis,gm.dim)
end

# TODO do i need this
function gradshape(gm::AbstractGaussianModel)
  return gm.dim*gm.nbasis
end

"""
Functions that compute the value of the basis functions
"""
function computebasis!(sys::AbstractSystem, gm::GaussianModel)
  k::Vector{Float64} = exp.( -0.5 .* (sys.x[1] .- gm.mus[1]).^2 ./ gm.sigs[1] )
  for i in 2:gm.dim
    k = vec( k * (exp.( -0.5 .* (sys.x[i] .- gm.mus[i]).^2 ./ gm.sigs[i] ))' )
  end
  gm.em .= k
end

function computebasis!(sys::AbstractSystem, gm::GaussianModelTD)
  k::Vector{Float64} = exp.( -0.5 .* (sys.x[1] .- gm.mus[1]).^2 ./ gm.sigs[1] )
  for i in 2:gm.dim
    k = vec( k * (exp.( -0.5 .* (sys.x[i] .- gm.mus[i]).^2 ./ gm.sigs[i] ))' )
  end
  k = vec( k * (exp.( -0.5 .* (sys.t .- gm.mus[end]).^2 ./ gm.sigs[end] ))' )
  gm.em .= k
end

"""
Force function with return value
"""
function force(sys::AbstractSystem, gm::AbstractGaussianModel; compbasis::Bool=true)
  if compbasis
    computebasis!(sys,gm)
  end
  fout::Vector{Float64} = zeros(Float64,gm.dim)
  for i in 1:gm.dim
    fout[i] = sum( @views gm.theta[:,i] .* gm.em )
  end
  return fout
end

"""
Force function with in-place changing
"""
function force!(sys::AbstractSystem, gm::AbstractGaussianModel; compbasis::Bool=true)
  if compbasis
    computebasis!(sys,gm)
  end
  for i in 1:gm.dim
    gm.f[i] = sum( @views gm.theta[:,i] .* gm.em )
  end
end

"""
Gradient wrt parameters for gaussian model
"""
function jacobian!(jac::Array{Float64},sys::AbstractSystem,gm::AbstractGaussianModel;
                   compbasis::Bool=false)
  if compbasis
    computebasis!(sys,gm)
  end
  ind0::Int64 = 1
  indf::Int64 = gm.nbasis
  for i in 1:gm.dim
    @. @views jac[ind0:indf,i] .= gm.em
    ind0 += gm.nbasis
    indf += gm.nbasis
  end
end

"""
Gradient wrt parameters for gaussian model
"""
function gradient!(grad::Array{Float64},sys::AbstractSystem,gm::AbstractGaussianModel;
                   compbasis::Bool=false)
  if compbasis
    computebasis!(sys,gm)
  end
  for i in 1:gm.dim
    grad[:,i] .= gm.em
  end
end

"""
Update parameters for the neural network model
"""
function updateparams!(lr::Float64,dtheta::Vector{Float64},gm::AbstractGaussianModel)
    # update coefficients for forces
    gm.theta .+= (lr .* reshape(dtheta, (gm.nbasis,gm.dim)))
end

"""
Set parameters for the neural network model
"""
function setparams!(theta::Vector{Float64},gm::AbstractGaussianModel)
    # update coefficients for forces
    gm.theta .= theta
end

function setparams!(gm::AbstractGaussianModel)
end

mutable struct GaussianModel2D <: AbstractSinglePotential
  dim::Int64
  nx::Int64
  ny::Int64
  nt::Int64
  mux::Vector{Float64}
  muy::Vector{Float64}
  mut::Vector{Float64}
  sigx::Vector{Float64}
  sigy::Vector{Float64}
  sigt::Vector{Float64}
  theta::Matrix{Float64}
  f::Vector{Float64}
  em::Vector{Float64}
  function GaussianModel2D(nx::Int64,ny::Int64,nt::Int64,
                           mux::Vector{Float64},muy::Vector{Float64},mut::Vector{Float64},
                           sigx::Vector{Float64},sigy::Vector{Float64},sigt::Vector{Float64},
                           theta::Matrix{Float64})
    new(2,nx,ny,nt,mux,muy,mut,sigx,sigy,sigt,theta,zeros(Float64,2),zeros(Float64,(nx*ny*nt)))
  end
end

function getdimensionality(gm::GaussianModel2D)
  return gm.dim
end

function jacshape(gm::GaussianModel2D)
  return (2*gm.nx*gm.ny*gm.nt,2)
end

function gradshape(gm::GaussianModel2D)
  return 2*gm.nx*gm.ny*gm.nt
end

"""
Force function with return value
"""
function force(sys::AbstractSystem, gm::GaussianModel2D)
  gm.em .= vec(kron(exp.(-0.5 .* (sys.t .- gm.mut).^2 ./ gm.sigt),
               kron(exp.(-0.5 .* (sys.x[2] .- gm.muy).^2 ./ gm.sigy),
                    exp.(-0.5 .* (sys.x[1] .- gm.mux).^2 ./ gm.sigx))))
  fout::Vector{Float64} = zeros(Float64,2)
  fout[1] = sum( @views gm.theta[:,1] .* gm.em )
  fout[2] = sum( @views gm.theta[:,2] .* gm.em )
  return fout
end

"""
Force function with in-place changing
"""
function force!(sys::AbstractSystem, gm::GaussianModel2D)
  gm.em .= vec(kron(exp.(-0.5 .* (sys.t .- gm.mut).^2 ./ gm.sigt),
               kron(exp.(-0.5 .* (sys.x[2] .- gm.muy).^2 ./ gm.sigy),
                    exp.(-0.5 .* (sys.x[1] .- gm.mux).^2 ./ gm.sigx))))
  gm.f[1] = sum( @views gm.theta[:,1] .* gm.em )
  gm.f[2] = sum( @views gm.theta[:,2] .* gm.em )
  #println(gm.f)
end

"""
Gradient wrt parameters for gaussian model
"""
function jacobian!(jac::Array{Float64},sys::AbstractSystem,gm::GaussianModel2D)
  @. @views jac[1:(gm.nx*gm.ny*gm.nt),1] .= gm.em
  @. @views jac[(gm.nx*gm.ny*gm.nt+1):(2*gm.nx*gm.ny*gm.nt),2] .= gm.em
end

"""
Gradient wrt parameters for gaussian model
"""
function gradient!(grad::Array{Float64},sys::AbstractSystem,gm::GaussianModel2D)
  for i in 1:gm.dim
    grad[:,i] .= gm.em
  end
  #grad[:,1] .= gm.em
  #grad[:,2] .= gm.em
end

"""
Update parameters for the neural network model
"""
function updateparams!(lr::Float64,dtheta::Vector{Float64},gm::GaussianModel2D)
    # update coefficients for forces
    gm.theta .+= (lr .* reshape(dtheta, (gm.nx*gm.ny*gm.nt,2)))
end

"""
Set parameters for the neural network model
"""
function setparams!(theta::Vector{Float64},gm::GaussianModel2D)
    # update coefficients for forces
    gm.theta .= theta
end

function setparams!(gm::GaussianModel2D)
end


text:
"""Implementation of Newton-with-restarted-Arnoldi propagation routine."""

using LinearAlgebra
using OffsetArrays

"""
```julia
    NewtonWrk(v0, m_max=10)
```

Workspace for the Newton-with-restarted-Arnoldi propagation routine.

Initializes the workspace for the propagation of a vector `v0`, using a
maximum Krylov dimension of `m_max` in each restart iteration. Note that
`m_max` should be smaller than the length of `v0`.
"""
mutable struct NewtonWrk{T}
    # arnoldi_vecs and v are pre-allocated for efficiency. Other values and
    # vectors are mainly for debugging, so that we can inspect internal
    # parameters like the Newton coefficients (a) after a call to newton!
    arnoldi_vecs :: Array{T}
    v :: T
    a :: OffsetVector{ComplexF64}
    leja :: OffsetVector{ComplexF64}
    radius :: Float64
    n_a :: Int64
    n_leja :: Int64
    restarts :: Int64
    function NewtonWrk(v0::T; m_max::Int64=10) where T
        if m_max >= length(v0)
            m_max = length(v0) - 1
        end
        new{T}(
            T[similar(v0) for _ in 1:m_max+1], # arnoldi_vecs
            similar(v0),                       # v
            OffsetVector(zeros(ComplexF64, 10*m_max+1), 0:10*m_max),  # a
            OffsetVector(zeros(ComplexF64, 10*m_max+1), 0:10*m_max),  # leja
            0.0,                               # radius
            0,                                 # n_a
            0,                                 # n_leja
            0,                                 # restarts
        )
    end
end

lbound(array::OffsetArray, dim=1) = first(axes(array)[dim])
ubound(array::OffsetArray, dim=1) = last(axes(array)[dim])

function leja_radius(z)
    r_max = maximum(map(abs, z))
    return 1.2 * r_max
end

"""
```julia
extend_leja!(leja, n, newpoints, n_use)
```

Given an array of `n` (ordered) Leja points, extract `n_use`
points from `newpoints`, and append them to the existing Leja points. The
array `leja` should be sufficiently large to hold the new Leja points,
which are appended after index `n_old`. It will be re-allocated if
necessary and may have a size of up to `2*(n+n_use)`.

# Arguments

- `leja`: Array of leja values. Must contain the "old" leja values to be kept
   in `leja(0:n-1)`. On output, `n_use` new leja points will be in
   `leja(n+:n+n_use-1)`, for the original value of `n`.  The `leja` array must
   use zero-based indexing.
- `n`: On input, number of "old" leja points in `leja`. On output, total number
  of leja points (i.e. `n=n+n_use`)
- `newpoints`: On input, candidate points for new leja points.  The `n_use`
  best values will be chosen and added to `leja`. On output, the values of
  `new_points` are undefined.
- `n_use`: Number of points that should be added to `leja`
"""
function extend_leja!(leja::OffsetVector{ComplexF64}, n,
                      newpoints::OffsetVector{ComplexF64}, n_use)
    @assert leja.offsets[1] == -1 "leja must be a zero-based vector"
    @assert newpoints.offsets[1] == -1 "newpoints must be a zero-based vector"
    if length(leja) < n + n_use
        # we allocate twice the required size, so we don't have to re-allocate
        # on every call
        resize!(leja, 2*(n + n_use))
        leja[n:end] .= 0
    end
    u = ubound(newpoints, 1)
    i_add_start = lbound(leja, 1)
    if n == 0
        # Use point of largest absolute value as starting point by moving it to
        # the end of the `newpoints` array
        z_last :: ComplexF64 = newpoints[u]
        for i = lbound(newpoints, 1) : (u-1)
            if abs(newpoints[i]) > abs(z_last)
                newpoints[u] = newpoints[i]
                newpoints[i] = z_last
                z_last = newpoints[u]
            end
        end
        leja[0] = newpoints[end]
        i_add_start = 1
    end
    exponent = 1.0 / (n + n_use)
    for i_add = i_add_start:n_use-1
        p_max :: Float64 = 0
        i_max = 0
        for i in lbound(newpoints, 1) : (u-i_add)
            p :: Float64 = 1
            for j = 0 : (n + i_add - 1)  # existing Leja points
                p = p * (abs(newpoints[i] - leja[j])^exponent)
            end
            if p > p_max
                p_max = p
                i_max = i
            end
        end
        # TODO: check that p_max did not overflow or underflow
        leja[n + i_add] = newpoints[i_max]
        # remove the used point by replacing it with the last unused point
        newpoints[i_max] = newpoints[u-i_add]
    end
    return n + n_use
end

"""
```julia
extend_newton_coeffs!(a, n_a, leja, func, n_leja, radius)
```

Extend the array `a` of existing Newton coefficients for the expansion of
the `func` from `n_a` coefficients to `n_leja` coefficients. Return a new value
`n_a=n_a+n_leja` with the total number of Newton coefficients in the updated
`a`.

# Arguments

- `a`: On input, a zero-based array of length `n_a` or greater, containing
  Newton coefficients. On output, array containing a total `n_leja`
  coefficients. The array `a` will be resized if necessary, and may have a
  length greater than `n_leja` on output
- `n_a`:  The number of Newton coefficients in `a`, on input. Elements of `a`
   beyond the first `n_a` elements will be overwritten.
- `leja`: Array of normalized Leja points, containing at least `n_leja`
  elements.
- `func`: Function for which to calcluate Newton coeffiecients
- `n_leja`: The number of elements in `leja` to use for calculating new
  coefficients, and the total number of Newton coefficients on output
- `radius`: Normalization radius for divided differences
"""
function extend_newton_coeffs!(a::OffsetVector{ComplexF64}, n_a::Int64,
                               leja::OffsetVector{ComplexF64}, func,
                               n_leja::Int64, radius::Float64)
    @assert a.offsets[1] == -1 "a must be a zero-based vector"
    m = n_leja - n_a
    n0 = n_a
    if length(a) < n_a + m
        # we allocate twice the required size, so we don't have to re-allocate
        # on every call
        resize!(a, 2*n_leja)
        a[n_a:end] .= 0
    end
    @assert length(leja) >= n_leja
    @assert radius > 0
    if (n_a == 0)
        a[0] = func(leja[0])
        n0 = 1
    end
    for k = n0 : (n_a - 1 + m)
        d :: ComplexF64 = 1
        pn :: ComplexF64 = 0
        for n = 1 : k-1
            zd = leja[k] - leja[n-1]
            d = d * zd / radius
            pn = pn + a[n] * d
        end
        zd = leja[k] - leja[k-1]
        d = d * zd / radius
        @assert abs(d) > 1e-200 "Divided differences too small"
        a[k] = (func(leja[k]) - a[0] - pn) / d
    end
    n_a = n_a + m
    return n_a
end

"""
```julia
newton!(Ψ, H, dt, wrk; func=(z -> exp(-1im*z)), norm_min=1e-14, relerr=1e-12,
        max_restarts=50)
```

Evaluate `Ψ = func(H*dt) Ψ` using a Newton-with-restarted-Arnoldi scheme.

# Arguments

- `Ψ`: The state to propagate, will be overwritten in-place with the propagated
  state
- `H`: Operator acting on `Ψ`. Together with `dt`, this is the argument to
  `func`
- `dt`: Implicit time step. Together with `H`, this is the argument to `func`
- `wkr`: Work array, initialized with [`NewtonWrk`](@ref)
- `func`: The function to apply to `H dt`, taking a single (scalar)
  complex-valued argument `z` in place of `H dt`. The default `func`
  is to evaluate the time evoluation operator for the Schrödinger equation
- `norm_min`: the minium norm at which to consider a state similar to `Ψ` as
  zero
- `relerr`: The relative error defining the convergence condition for the
  restart iteration. Propagation stops when the norm of the accumulated `Ψ`
  is stable up to the given relative error
- `max_restart`: The maximum number of restart iterations. Exceeding
  `max_restart` will throw an `AssertionError`.
"""
function newton!(Ψ, H, dt, wrk; kwargs...)
    func = get(kwargs, :func, z -> exp(-1im * z))
    norm_min :: Float64 = get(kwargs, :norm_min, 1e-14)
    relerr :: Float64 = get(kwargs, :relerr, 1e-12)
    max_restarts :: Int64 = get(kwargs, :max_restarts, 50)

    m_max :: Int64 = length(wrk.arnoldi_vecs) - 1
    m :: Int64 = m_max
    fill!(wrk.a, 0)
    fill!(wrk.leja, 0)
    R = Array{ComplexF64,1}(undef, m+1)
    P = Array{ComplexF64,1}(undef, m+1)
    R_abs = Array{Float64,1}(undef, m+1)
    Hess = zeros(ComplexF64, m_max+1, m_max+1)
    # there seems to be a problem declaring dt as Float64 directly (if you pass
    # it an Int, the kwargs seems to prevent Julia from auto-converting it
    _dt :: Float64 = dt

    @assert _dt ≠ 0.0

    n_a :: Int64 = 0  # number of Newton coefficients (accumulated)
    n_leja :: Int64 = 0  # number of Leja points (accumulated)
    copyto!(wrk.v, Ψ)

    s = 0  # counter for the restart loop
    β :: Float64 = norm(wrk.v)
    lmul!(1/β, wrk.v)  # normalize

    while true  # restart loop (s → s+1)

        m = arnoldi!(Hess, wrk.arnoldi_vecs, m, wrk.v, H, _dt;
                     extended=true, norm_min=norm_min)
        if m == 1 && s == 0
            λ = β * Hess[1, 1]
            # wrk.v is an eigenvec of H with eigenval Hess[1, 1],
            # so Ψ is an eigenvec with eigenval λ, and f(H) Ψ = f(λ) Ψ
            lmul!(func(λ), Ψ)
            break
        end
        ritz = diagonalize_hessenberg_matrix(Hess, m, accumulate=true)

        # In the first iteration, the radius will be determined
        if s == 0
            wrk.radius = leja_radius(ritz)
        end

        # Get the Leja points (i.e. Ritz values in the proper order
        n_s = n_leja  # we need to keep the old n_leja
        n_leja = extend_leja!(
            wrk.leja, n_leja, OffsetVector(ritz, 0:length(ritz)-1), m
        )

        # Extend the Newton coefficients
        n_a = extend_newton_coeffs!(
            wrk.a, n_a, wrk.leja, func, n_leja, wrk.radius
        )
        @assert n_a == n_leja

        # allocate array R for Newton basis polynomicals in the extended
        # Hessenberg matrix and P for the full Newton series
        if length(R) ≠ m+1
            # we always treat R, R_abs, and P together
            resize!(R, m+1)
            resize!(R_abs, m+1)
            resize!(P, m+1)
        end

        # Evaluate Newton Polynomial in the extended Hessenberg matrix
        P .= 0.0
        R .= 0.0
        R[1] = β
        P[1] = wrk.a[n_s] * β
        for k = 1:m-1
            R[1:m+1] = (
                (Hess[1:m+1,1:m+1] * R[1:m+1] - wrk.leja[n_s+k-1] * R[1:m+1])
                / wrk.radius
            )
            # P += a_{n_s+k} R
            axpy!(wrk.a[n_s+k], view(R, 1:m+1), view(P, 1:m+1))
        end

        # Calculate the new solution
        if s == 0
            fill!(Ψ, 0)
        end
        # Ψ += P[i] * arnoldi_vecs[i]
        for i = 1:m
          axpy!(P[i], wrk.arnoldi_vecs[i], Ψ)
        end

        # Calculate the starting vector v_{s+1} for the next iteration
        R[1:m+1] = (
            ((Hess[1:m+1,1:m+1] * R[1:m+1]) - wrk.leja[n_s+m-1] * R[1:m+1])
            / wrk.radius
        )
        R_abs[1:m+1] = abs.(R[1:m+1])
        β = norm(R_abs[1:m+1])
        lmul!(1/β, view(R, 1:m+1))
        copyto!(wrk.arnoldi_vecs[1], wrk.v)
        lmul!(R[1], wrk.v)
        for i = 2 : m+1
            axpy!(R[i], wrk.arnoldi_vecs[i], wrk.v)
        end

        # Convergence check (relative error of Newton series)
        Ψ_relerr = β * abs(wrk.a[n_a-1]) / (1 + norm(Ψ))
        if Ψ_relerr < relerr
            break
        else
            s += 1
            @assert s <= max_restarts
        end

    end

    wrk.restarts = max(0, s - 1)
    wrk.n_leja = n_leja
    wrk.n_a = n_a
    return Ψ

end


text:
"""

```
icell3D
```

Returns the i,j,k coordinates of the cell from the coordinates of an atom and box properties

"""
function icell3D(x::AbstractVector, box::Box)
    i = trunc(Int, (x[1] + box.sides[1] / 2) / box.l[1]) + 1
    j = trunc(Int, (x[2] + box.sides[2] / 2) / box.l[2]) + 1
    k = trunc(Int, (x[3] + box.sides[3] / 2) / box.l[3]) + 1
    return i, j, k
end

#
# given the 1D representation of the index of a 3D cell, returns the 3 indexes of the
# coordinates in 3 dimensions
#

# if the number of cells in each dimension is different

function icell3D(n::AbstractVector{Int}, i1D)
    jk = i1D % (n[2] * n[3])
    if jk == 0
        jk = n[2] * n[3]
    end
    j, k = icell2D(n[2:3], jk)
    i = round(Int, (i1D - jk) / (n[2] * n[3])) + 1
    return i, j, k
end

# If the number of cells in each dimension is the same

function icell3D(n::Int, i1D::Int)
    n2 = n^2
    jk = i1D % n2
    if jk == 0
        jk = n2
    end
    j, k = icell2D(n, jk)
    i = round(Int, (i1D - jk) / (n2)) + 1
    return i, j, k
end

# Given the 1D represetation of the index of a 2D cell, returns the 2
# indexes of the 2D representation

# If the two dimensions are different

function icell2D(n::AbstractVector{Int}, i1D::Int)
    j = i1D % n[2]
    if j == 0
        j = n[2]
    end
    i = round(Int, (i1D - j) / n[2]) + 1
    return i, j
end

# If the two dimensions are the same

function icell2D(n::Int, i1D::Int)
    j = i1D % n
    if j == 0
        j = n
    end
    i = round(Int, (i1D - j) / n) + 1
    return i, j
end


text:
################################################################################
#
#	CONCRETE TYPE
#
#   ReciprocalUnitcell{P,B} <: AbstractReciprocalUnitcell{P,B}
#   --> P is the reciprocal site type (<: AbstractReciprocalPoint{D})
#       --> D is the dimension of embedding space
#   --> B is the bond type (<: AbstractBond{L,N})
#       --> N is the dimension of the Bravais lattice that the bond is located in
#       --> L is the label type of bonds
#
#   FILE CONTAINS
#       - abstract type definition
#       - interface definition
#       - TODO interface testing
#
################################################################################

################################################################################
#
#   ABSTRACT TYPE DEFINITION
#
################################################################################
mutable struct ReciprocalUnitcell{P,B} <: AbstractReciprocalUnitcell{P,B}

    # basis vectors of the Bravais lattice (of the reciprocal unitcell)
    lattice_vectors	:: Vector{Vector{Float64}}

    # the gamma point (only site within the reciprocal unitcell)
    gamma_point     :: P

    # list of bonds
    bonds			:: Vector{B}

end

# export the type
export ReciprocalUnitcell

################################################################################
#
#	INTERFACING / ACCESSING RECIPROCAL UNITCELLS
#	(functions that have to be overwritten by concrete types)
#
################################################################################

# constructor
function newReciprocalUnitcell(
            ::Type{ReciprocalUnitcell{P,B}},
            lattice_vectors :: Vector{<:Vector{<:Real}},
            bonds           :: Vector{B}
        ) :: ReciprocalUnitcell{P,B} where {D,N,L,P<:AbstractReciprocalPoint{D},B<:AbstractBond{L,N}}

    # call the constructor
    return ReciprocalUnitcell{P,B}(
        lattice_vectors,
        newReciprocalPoint(P, zeros(D), "Gamma", "\\Gamma"),
        bonds
    )
end

# export the constructor
export newReciprocalUnitcell

# accessing a list of lattice vectors
function latticeVectors(
            unitcell :: ReciprocalUnitcell{P,B}
        ) :: Vector{Vector{Float64}} where {D,N,L,P<:AbstractReciprocalPoint{D},B<:AbstractBond{L,N}}

    # return the lattice vectors
    return unitcell.lattice_vectors
end
# setting a list of lattice vectors
function latticeVectors!(
            unitcell        :: ReciprocalUnitcell{P,B},
            lattice_vectors :: Vector{<:Vector{<:Real}}
        ) where {D,N,L,P<:AbstractReciprocalPoint{D},B<:AbstractBond{L,N}}

    # set the lattice vectors
    unitcell.lattice_vectors = lattice_vectors
end

# export the latticeVectors interface
export latticeVectors, latticeVectors!

# accessing the gamma point
function gammaPoint(
            unitcell :: ReciprocalUnitcell{P,B}
        ) :: P where {D,N,L,P<:AbstractReciprocalPoint{D},B<:AbstractBond{L,N}}

    # return the gamma point
    return unitcell.gamma_point
end

# setting the gamma point
function gammaPoint!(
            unitcell :: ReciprocalUnitcell{P,B},
            site     :: P
        ) where {D,N,L,P<:AbstractReciprocalPoint{D},B<:AbstractBond{L,N}}

    # set the gamma point
    unitcell.gamma_point = site
end

# export gamma point interface
export gammaPoint, gammaPoint!

# accessing a list of bonds
function bonds(
            unitcell :: ReciprocalUnitcell{P,B}
        ) :: Vector{B} where {D,N,L,P<:AbstractReciprocalPoint{D},B<:AbstractBond{L,N}}

    # return the internal list of bonds
    return unitcell.bonds
end
# setting a list of bonds
function bonds!(
            unitcell :: ReciprocalUnitcell{P,B},
            bonds    :: Vector{B}
        ) where {D,N,L,P<:AbstractReciprocalPoint{D},B<:AbstractBond{L,N}}

    # set the internal list of bonds
    unitcell.bonds = bonds
end

# export the bonds interface
export bonds, bonds!


text:
struct PointMass3DMissile <: AbstractMissile
end

function State(env::PointMass3DMissile)
    return function (p=zeros(3), v=zeros(3))
        ComponentArray(p=p, v=v)
    end
end

"""
u ∈ ℝ³: acceleration
"""
function Dynamics!(env::PointMass3DMissile)
    @Loggable function dynamics!(dx, x, params, t; u)
        @unpack p, v = x
        @log p
        @log v
        dx.p = v
        dx.v = u
    end
end