text:
# This file must be loaded _after_ raysplitting

export evolve!, evolve, timeseries!, timeseries

"""
    evolve!([p::AbstractParticle,] bd::Billiard, t)
Evolve the given particle `p` inside the billiard `bd`. If `t` is of type
`AbstractFloat`, evolve for as much time as `t`. If however `t` is of type `Int`,
evolve for as many collisions as `t`.
Return the states of the particle between collisions.

This function mutates the particle, use `evolve` otherwise. If a particle is
not given, a random one is picked through [`randominside`](@ref).

### Return

* `ct::Vector{T}` : Collision times.
* `poss::Vector{SVector{2,T}}` : Positions at the collisions.
* `vels::Vector{SVector{2,T}})` : Velocities exactly after the collisions.
* `ω`, either `T` or `Vector{T}` : Angular velocity/ies (returned only for magnetic
  particles).

The time `ct[i+1]` is the time necessary to reach state `poss[i+1], vels[i+1]` starting
from the state `poss[i], vels[i]`. That is why `ct[1]` is always 0 since
`poss[1], vels[1]` are the initial conditions. The angular velocity `ω[i]` is the one
the particle has while propagating from state `poss[i], vels[i]` to `i+1`.

Notice that at any point, the velocity vector `vels[i]` is the one obdained *after*
the specular reflection of the `i-1`th collision.

### Ray-splitting billiards
    evolve!(p, bd, t, raysplitters)

To implement ray-splitting, the `evolve!` function is supplemented with a
fourth argument, `raysplitters` which is a tuple of [`RaySplitter`](@ref) instances.
Notice that `evolve` **always mutates the billiard** if ray-splitting is used!
For more information and instructions on using ray-splitting
please visit the official documentation.
"""
function evolve!(p::AbstractParticle{T}, bd::Billiard{T}, t, raysplitters = nothing;
    warning = false) where {T<:AbstractFloat}

    if t ≤ 0
        throw(ArgumentError("`evolve!()` cannot evolve backwards in time."))
    end
    if ispinned(p, bd)
        push!(rpos, p.pos)
        push!(rvel, p.vel)
        push!(rt, Inf)
        return (rt, rpos, rvel, p.ω)
    end

    ismagnetic = p isa MagneticParticle
    isray = !isa(raysplitters, Nothing)
    isray && acceptable_raysplitter(raysplitters, bd)
    raysidx = raysplit_indices(bd, raysplitters)
    ismagnetic && isray && (omegas = [p.ω])

    rt = T[0.0]; rpos = [p.pos]; rvel = [p.vel]
    count = zero(t)
    flight = zero(T)
    if typeof(t) == Int
        for zzz in (rt, rpos, rvel)
            sizehint!(zzz, t)
        end
    end

    while count < t

        i, tmin, pos, vel = bounce!(p, bd, raysidx, raysplitters)
        flight += tmin

        if isperiodic(i, bd)
            continue
        else
            push!(rpos, pos + p.current_cell)
            push!(rvel, vel)
            push!(rt, flight)
            ismagnetic && isray && push!(omegas, p.ω)
            # set counter
            count += increment_counter(t, flight)
            flight = zero(T)
        end
    end#time, or collision number, loop

    # Return stuff
    if ismagnetic && isray
        return (rt, rpos, rvel, omegas)
    elseif ismagnetic
        return (rt, rpos, rvel, p.ω)
    else
        return (rt, rpos, rvel)
    end
end

"""
    evolve(p, args...)
Same as [`evolve!`](@ref) but copies the particle instead.
"""
evolve(p::AbstractParticle, args...) = evolve!(copy(p), args...)
evolve(bd::Billiard, args...; kwargs...) =
    evolve!(randominside(bd), bd, args...; kwargs...)

#####################################################################################
# Timeseries
#####################################################################################
"""
    timeseries!([p::AbstractParticle,] bd::Billiard, t; dt, warning)
Evolve the given particle `p` inside the billiard `bd` for the condition `t`
and return the x, y, vx, vy timeseries and the time vector.
If `t` is of type `AbstractFloat`, then evolve for as much time as `t`.
If however `t` is of type `Int`, evolve for as many collisions as `t`.
Otherwise, `t` can be any function, that takes as an input `t(n, τ, i, p)`
and returns `true` when the evolution should terminate. Here `n` is the amount
of obstacles collided with so far, `τ` the amount time evolved so far, `i`
the obstacle just collided with and `p` the particle (so you can access e.g.
`p.pos`).

This function mutates the particle, use `timeseries` otherwise. If a particle is
not given, a random one is picked through [`randominside`](@ref).

The keyword argument `dt` is the time step used for interpolating the time
series in between collisions. `dt` is capped by the collision time, as
the interpolation _always_ stops at collisions.
For straight propagation `dt = Inf`, while for magnetic `dt = 0.01`.

For pinned magnetic particles, `timeseries!` issues a warning and returns the
trajectory of the particle. If `t` is integer, the trajectory is evolved for
one full circle only.

Internally uses [`DynamicalBilliards.extrapolate`](@ref).

### Ray-splitting billiards
    timeseries!(p, bd, t, raysplitters; ...)

To implement ray-splitting, the `timeseries!` function is supplemented with a
fourth argument, `raysplitters` which is a tuple of [`RaySplitter`](@ref) instances.
Notice that `timeseries` **always mutates the billiard** if ray-splitting is used!
For more information and instructions on using ray-splitting
please visit the official documentation.
"""
function timeseries!(p::AbstractParticle{T}, bd::Billiard{T}, f, raysplitters = nothing;
                     dt = typeof(p) <: Particle ? T(Inf) : T(0.01),
                     warning::Bool = true) where {T}

    ts = [zero(T)]
    x  = [p.pos[1]]; y  = [p.pos[2]]
    vx = [p.vel[1]]; vy = [p.vel[2]]
    ismagnetic = p isa MagneticParticle
    prevω = ismagnetic ? p.ω : T(0)
    isray = !isa(raysplitters, Nothing)
    isray && acceptable_raysplitter(raysplitters, bd)
    raysidx = raysplit_indices(bd, raysplitters)
    n, i, t, flight = 0, 0, zero(T), zero(T)
    prevpos = p.pos + p.current_cell
    prevvel = p.vel

    if ispinned(p, bd)
        warning && @warn "Pinned particle detected!"
        #return one cycle if t is a collision number
        t_ret = typeof(t) <: Integer ? 2π/p.ω : T(t)
        nx, ny, nvx, nvy, nts = extrapolate(p, prevpos, prevvel, t_ret, dt, p.ω)
        append!(ts, nts[2:end] .+ t)
        append!(x, nx[2:end]); append!(vx, nvx[2:end])
        append!(y, ny[2:end]); append!(vy, nvy[2:end])
        return x, y, vx, vy, ts
    end

    @inbounds while check_condition(f, n, t, i, p) # count < t
        i, ct = bounce!(p, bd, raysidx, raysplitters)
        flight += ct
        if isperiodic(i, bd) # do nothing at periodic obstacles
            continue
        else
            if flight ≤ dt # push collision point only
                push!(ts, flight + t)
                push!(x, p.pos[1] + p.current_cell[1])
                push!(y, p.pos[2] + p.current_cell[2])
                push!(vx, p.vel[1]); push!(vy, p.vel[2])
            else
                nx, ny, nvx, nvy, nts = extrapolate(p, prevpos, prevvel, flight, dt, prevω)
                append!(ts, nts[2:end] .+ t)
                append!(x, nx[2:end]); append!(vx, nvx[2:end])
                append!(y, ny[2:end]); append!(vy, nvy[2:end])
            end
            prevpos = p.pos + p.current_cell
            prevvel = p.vel
            t += flight; n += 1
            flight = zero(T)
            ismagnetic && isray && (prevω = p.ω)
        end
    end
    return x, y, vx, vy, ts
end

check_condition(f::Int, n, t, i, p) = f > n
check_condition(f::AbstractFloat, n, t, i, p) = f > t
check_condition(f, n, t, i, p) = !f(n, t, i, p)

"""
    extrapolate(particle, prevpos, prevvel, ct, dt[, ω]) → x, y, vx, vy, t

Create the timeseries that connect a `particle`'s previous position and velocity
`prevpos, prevvel` with the `particle`'s current position and velocity,
provided that the collision time between previous and current state is `ct`.

`dt` is the sampling time and if the `particle` is `MagneticParticle` then
you should provide `ω`, the angular velocity that governed the free flight.

Here is how this function is used (for example)
```julia
prevpos, prevvel = p.pos + p.current_cell, p.vel
for _ in 1:5
    i, ct, pos, vel = bounce!(p, bd)
    x, y, x, vy, t = extrapolate(p, prevpos, prevvel, ct, 0.1)
    # append x, y, ... to other vectors or do whatever with them
    prevpos, prevvel = p.pos + p.current_cell, p.vel
end
```
"""
function extrapolate(p::MagneticParticle{T}, prevpos::SV{T}, prevvel::SV{T}, ct::T,
                     dt::T, ω::T = p.ω) where {T <: AbstractFloat}

    φ0 = atan(prevvel[2], prevvel[1])
    s0, c0 = sincos(φ0)

    tvec = collect(0:dt:ct)
    x = Vector{T}(undef, length(tvec))
    y = Vector{T}(undef, length(tvec))
    vx = Vector{T}(undef, length(tvec))
    vy = Vector{T}(undef, length(tvec))

    @inbounds for (i,t) ∈ enumerate(tvec)
        s,c = sincos(ω*t + φ0)
        x[i] = s/ω + prevpos[1] - s0/ω
        y[i] = -c/ω + prevpos[2] + c0/ω
        vx[i] = c; vy[i] = s
    end

    # finish with ct
    @inbounds if tvec[end] != ct
        push!(tvec, ct)
        push!(x, p.pos[1] + p.current_cell[1])
        push!(y, p.pos[2] + p.current_cell[2])
        push!(vx, p.vel[1]); push!(vy, p.vel[2])
    end

    return x, y, vx, vy, tvec
end

extrapolate(p::Particle, a, b, c, d, ω) = extrapolate(p::Particle, a, b, c, d)
function extrapolate(p::Particle{T}, prevpos::SV{T}, prevvel::SV{T}, ct::T,
                     dt::T) where {T <: AbstractFloat}

    tvec = collect(0:dt:ct)
    x = Vector{T}(undef, length(tvec))
    y = Vector{T}(undef, length(tvec))
    vx = Vector{T}(undef, length(tvec))
    vy = Vector{T}(undef, length(tvec))

    @inbounds for (i,t) ∈ enumerate(tvec)
        x[i] = prevpos[1] + t*prevvel[1]
        y[i] = prevpos[2] + t*prevvel[2]
        vx[i], vy[i] = prevvel
    end

    # finish with ct
    @inbounds if tvec[end] != ct
        push!(tvec, ct)
        push!(x, p.pos[1] + p.current_cell[1])
        push!(y, p.pos[2] + p.current_cell[2])
        push!(vx, p.vel[1]); push!(vy, p.vel[2])
    end

    return x, y, vx, vy, tvec
end

# non-mutating version
"""
    timeseries(p, args...; kwargs...)
Non-mutating version of [`timeseries!`](@ref)
"""
@inline timeseries(p::AbstractParticle, args...; kwargs...) =
    timeseries!(copy(p), args...; kwargs...)


text:
"""
Types and definitions used for simulations of the uniform electron gas.
"""
module UniformElectronGas

using ..CPIMC
using ..CPIMC.PlaneWaves

import ..CPIMC: energy, w

import StaticArrays: StaticVector
import LinearAlgebra: dot

export UEG, β, EF, λ, β_Ha, rs, kF

"""
Singleton for dispatching.

For other models this would be the place to store additional
parameters that might come up in the matrix elements.
"""
struct UEG <: Model end

@doc raw"""
    energy(m::UEG, o::PlaneWave)

The single-particle energy of a plane wave with momentum $\vec{k}$ is given by (Hartree atomic units):

$ \epsilon_k = \vec{k}^2 / 2. $

In the internal unit system we use integer $k$-values, which results in:

$ \epsilon_k = \vec{k}^2. $

"""
function energy(m::UEG, o::PlaneWave)
    dot(o.vec,o.vec)
end

"""
    w(i::Orbital, j::Orbital, k::Orbital, l::Orbital)

Return the two-particle matrix element of the Coulomb interaction.
Zero is returned if momentum or spin is not conserved, in consequence of the Bloch-theorem
and the spin kronecker-delta in the plane-spin-wave basis.

The singular contribution `(i.vec == k.vec)` is also removed, i.e. set to zero.
This property is not a result of the Bloch theorem for the two-particle Coulomb matrix
element in the plane wave basis but is added here to allow for straightforward summation over
all combination of orbitals without explictly excluding this component.

It is not part of the UEG many-body Hamiltonian due to the physical assumption that contributions from the
uniform background cancel with this diverging term in the thermodynamic limit.
"""
function w(m::UEG, i::Orbital, j::Orbital, k::Orbital, l::Orbital)
    if !iszero(i.vec + j.vec - k.vec - l.vec) | (i.spin != k.spin) | (j.spin != l.spin)# momentum and spin conservation
        return 0.0
    elseif i.vec == k.vec
        # here return 0 in order to remove this term from sums over all occupations since it is canceled by the uniform background in the TD-Limes
        0.0
    else
        return 1.0 / dot(i.vec - k.vec, i.vec - k.vec)
    end
end

"""
    λ(N, rs, d::Int)

calculate λ, which is the coupling constant of the Hamiltonian in internal units
from the particle number N, Bruecker parameter rs for the density and spatial dimension d
the differences for d ∈ {2,3} reflect the different expressions for the Coulomb interaction, whichs Fourier component is v_q = \frac{4π}{q^2} in 3D and v_q = \frac{2π}{q} in 2D
the one-dimensional case involves a short-distance cutoff to avoid the non-integrable divergence of the interaction at zero distance and is not implemented
"""
function λ(N, rs, d::Int)
    if d == 3
        (4 / (2π)^3 ) * cbrt(4π*N/3) * rs
    elseif d == 2
        sqrt(N/π) * rs / (2π)
    else
        throw(DomainError("coupling constant λ is only implemented for dimension d ∈ {2,3} and not for d=$(d)"))
    end
end

"""
    rs(N, λ, d::Int)
calculate the Brueckner parameter rs for the density from the particle number N
and λ, which is the coupling constant of the Hamiltonian in internal units"""
function rs(N, λ, d::Int)
    if d == 3
        ( (2π)^3 / 4 ) * cbrt(3 / (4π*N)) * λ
    elseif d == 2
        2π * sqrt(π/N) * λ
    else
        throw(DomainError("Coupling constant λ is only implemented for dimension d ∈ {2,3} and not for d=$(d). Thus, no conversion from λ to rs is defined for d ∉ {2,3}"))
    end
end

"""
    internal_energy_factor(N, rs, d)

This factor is used to convert an energy quantity from Hartree a.u. to internal units.
A value of the quantity in Hartree a.u. has to be multiplied by this factor to yield the corresponding value of the quantity in internal units.
A value of the quantity in internal units has to be divided by this factor to yield the corresponding value of the quantity in Hartree a.u.
"""
internal_energy_factor(N, rs, d) = (2π)^4 * λ(N, rs, d)^2 / 8

"""
    kF(rs, ξ, d::Int)
calculate the Fermi wavenumber in Hartree a.u. from Brueckner parameter rs for the density,
the fractional spin polarization ξ and spatial dimension d
"""
function kF(rs, ξ, d::Int)
    if d == 1
        α = 4 / π
    elseif d == 2
        α = 1 / sqrt(2)
    elseif d == 3
        α = cbrt( 4 / (9π) )
    else
        throw(DomainError("Fermi-wavenumber only defined for 1-, 2- and 3-dimensional systems. Choose d ∈ {1,2,3} and not d=$(d)"))
    end
    (1 + ξ)^(1/d) / ( α*rs )
end

"""
    EF(rs, ξ, d::Int)
calculate the Fermi energy in Hartree a.u. from Brueckner parameter rs for the density,
the fractional spin polarization ξ and spatial dimension d
"""
EF(rs, ξ, d::Int) = kF(rs, ξ, d)^2 / 2

"""
    β_Ha(Θ, EF)
calculate the inverse temperature in Hartree a.u. from the reduced temperature Θ and the Fermi energy EF in Hartree a.u.
the inverse temperature is defined by β = 1 / (kB*T) where kB is the Boltzmann constant and T is the temperature
"""
β_Ha(Θ, EF) = 1 / (Θ * EF)

"""
    β(Θ, rs, N, ξ, d::Int)

calculate β in internal units
"""
function β(Θ, N, ξ, d::Int)
    #calculate the factor alpha used in calculation of the fermi-vector in Ha units
    if d == 1
        α = 4 / π
    elseif d == 2
        α = 1 / sqrt(2)
    elseif d == 3
        α = cbrt( 4 / (9π) )
    else
        throw(DomainError("Fermi-wavenumber only defined for 1-, 2- and 3-dimensional systems. Choose d ∈ {1,2,3} and not d=$(d)"))
    end

    #calculate λ_rs_ratio
    if d == 3
        λ_rs_ratio = (4 / (2π)^3 ) * cbrt(4π*N/3)
    elseif d == 2
        λ_rs_ratio = sqrt(N/π) / (2π)
    else
        throw(DomainError("coupling constant λ is only implemented for dimension d ∈ {2,3} and not for d=$(d)"))
    end

    return (16/Θ) * (α / ((1 + ξ)^(1/d) * (2π)^2 * λ_rs_ratio))^2
end

end

ext: jl


text:
using JSON
using Random
using CUDA
using Statistics
using Dates
using DataStructures

include("environments/collect_points.jl")
include("environments/dummy_app.jl")
include("brains/continuous_time_rnn.jl")
include("optimizers/cma_es.jl")
include("optimizers/randomizer.jl")
include("tools/episode_runner.jl")
include("tools/write_results.jl")

struct TrainingCfg
    number_generations::Int
    number_validation_runs::Int
    number_rounds::Int
    maximum_env_seed::Int
    environment::OrderedDict
    brain::OrderedDict
    optimizer::OrderedDict
    experiment_id::Int

    function TrainingCfg(configuration::OrderedDict)
        new(
            configuration["number_generations"],
            configuration["number_validation_runs"],
            configuration["number_rounds"],
            configuration["maximum_env_seed"],
            configuration["environment"],
            configuration["brain"],
            configuration["optimizer"],
            get(configuration, "experiment_id", -1),
        )
    end
end

function main()

    configuration_file = "configurations/CMA_ES_Deap_CTRNN_Dense_Dummy_App.json"

    # Load configuration file
    configuration = JSON.parsefile(configuration_file, dicttype = OrderedDict)

    config = TrainingCfg(configuration)

    # Get environment type from configuration
    # TODO: Choose environment type from configuration
    environment_type = DummyApp

    # Get brain type from configuration
    # TODO: Choose brain type from configuration
    brain_type = ContinuousTimeRNN

    # Get optimizer type from configuration
    if configuration["optimizer"]["type"] == "Randomizer"
        optimizer_type = OptimizerRandom
    elseif configuration["optimizer"]["type"] == "CMA-ES-Deap"
        optimizer_type = OptimizerCmaEs
    else
        error("No valid optimizer type")
    end

    number_individuals = config.optimizer["population_size"]

    # Initialize environments
    environments = environment_type(config.environment, number_individuals)

    number_inputs = get_number_inputs(environments)
    number_outputs = get_number_outputs(environments)

    # Initialize brains
    brains = brain_type(config.brain, number_inputs, number_outputs, number_individuals)

    individual_size = get_individual_size(brains)

    # Initialize optimizer
    optimizer = optimizer_type(individual_size, config.optimizer)

    best_genome_overall = nothing
    best_reward_overall = typemin(Int32)

    required_shared_memory = get_memory_requirements(brains) + get_memory_requirements(environments)

    # Get start time of training and date
    start_time_training = now()

    log = OrderedDict()

    # Run evolutionary training for given number of generations
    for generation = 1:config.number_generations

        start_time_generation = time()

        # Environment seed for this generation (excludes validation environment seeds)
        env_seed = Random.rand(config.number_validation_runs:config.maximum_env_seed)

        # Ask optimizer for new population
        individuals = ask(optimizer)

        # Training runs for current generation
        rewards_training = training_runs(
            individuals,
            number_individuals,
            brains,
            environments,
            env_seed = env_seed,
            number_rounds = config.number_rounds,
            threads = brains.number_neurons,
            blocks = number_individuals,
            shared_memory = required_shared_memory,
        )

        # Tell optimizer new rewards
        tell(optimizer, rewards_training)

        best_genome_current_generation = individuals[(findmax(rewards_training))[2], :]

        # Validation runs for best individual in current generation
        rewards_validation = validation_runs(
            best_genome_current_generation,
            individual_size,
            config.number_validation_runs,
            brains,
            environments,
            threads = brains.number_neurons,
            blocks = config.number_validation_runs,
            shared_memory = required_shared_memory,
        )

        best_reward_current_generation = mean(rewards_validation)

        # Better individual found in current generation?
        if best_reward_current_generation > best_reward_overall
            best_genome_overall = best_genome_current_generation
            best_reward_overall = best_reward_current_generation
        end

        # Logging and printing
        elapsed_time_current_generation = round(time() - start_time_generation, digits = 3)
        log_line = OrderedDict()
        log_line["gen"] = generation
        log_line["min"] = minimum(rewards_training)
        log_line["mean"] = mean(rewards_training)
        log_line["max"] = maximum(rewards_training)
        log_line["best"] = best_reward_overall
        log_line["elapsed_time"] = elapsed_time_current_generation
        log[generation] = log_line
        println(
            "Generation: ", generation,
            "   Min: ", minimum(rewards_training),
            "   Mean: ", mean(rewards_training),
            "   Max: ", maximum(rewards_training),
            "   Best: ", best_reward_overall,
            "   Elapsed time: ", elapsed_time_current_generation, "s"
        )
    end

    result_directory =
        "Simulation_results/" *
        Dates.format(floor(start_time_training, Second), DateFormat("yyyy-mm-dd_HH-MM-SS"))

    mkdir(result_directory)
    write_results_to_textfile(
        result_directory * "/Log.txt",
        configuration,
        log,
        number_inputs,
        number_outputs,
        number_individuals,
        individual_size,
        now() - start_time_training,
    )
end

main()


text:
using Distributions

function rCIR(n::Integer, Dt::Real, x0::Real, δ, γ, σ)
    β = γ/σ^2*exp(2*γ*Dt)/(exp(2*γ*Dt)-1)
    if n == 1
        ks = rand(Poisson(γ/σ^2*x0/(exp(2*γ*Dt)-1)))
        return rand(Gamma(ks+δ/2, 1/β))
    else
        ks = rand(Poisson(γ/σ^2*x0/(exp(2*γ*Dt)-1)), n)
        return rand.(Gamma.(ks .+ δ/2, 1/β))
    end
end

transition_CIR(n::Integer, Dt::Real, x0::Real, δ, γ, σ) = rCIR(n::Integer, Dt::Real, x0::Real, δ, γ, σ)

function rec_transition_CIR(Dts, x, δ, γ, σ)
    x_new = transition_CIR(1, Dts[1], x[end], δ, γ, σ)
    if length(Dts) == 1
        return Float64[x; x_new]
    else
        return Float64[x; rec_transition_CIR(Dts[2:end], x_new, δ, γ, σ)]
    end
end

function generate_CIR_trajectory(times, x0, δ, γ, σ)
    Dts = diff(times)
    return rec_transition_CIR(Dts, [x0], δ, γ, σ)
end

function next_wms_from_wms_prime(wms_prime, Λ_prime, y, θ_prime, α)#update
    #Make sure we deal correctly with weights equal to 0
    #Probably harmonise what can be improved in the filtering algorithm
    unnormalised_wms = wms_prime .* DualOptimalFiltering.μπh(Λ_prime, θ_prime, α, y)
    return unnormalised_wms |> DualOptimalFiltering.normalise
end

μπh(m, θ, α, y; λ = 1) = exp.(logμπh(m, θ, α, y; λ = λ))

function logμπh_inside(m, θ, α, y; λ = 1)
    return map(mm -> logμπh(mm::Integer, θ, α, y; λ = λ), m)
end

function logμπh(m::AbstractArray{U, 1}, θ, α, y; λ = 1)  where U<:Integer
    return logμπh_inside(m, θ, α, y; λ = λ)
end

function logμπh(m::Integer, θ, α, y::Integer; λ = 1)
    return y*log(λ) + (α+m)*log(θ) + lgamma_local(m+y+α) - SpecialFunctions.logfactorial(y) - lgamma_local(α+m) - (y+α+m) * log(θ + λ)
end
function logμπh(m::Integer, θ, α, y::AbstractArray{T, 1}; λ = 1) where T <: Integer
    s = sum(y)
    n = length(y)
    return s*log(λ) + (α+m)*log(θ) + lgamma_local(m+s+α) - sum(SpecialFunctions.logfactorial.(y)) - lgamma_local(α+m) - (s+α+m) * log(θ + n*λ)
end

function next_wms_prime_from_wms(wms, Λ, Δt, θ, γ, σ)
    nΛ = length(Λ)
    wms_prime = zeros(maximum(Λ)+1)
    p = γ/σ^2*1/(θ*exp(2*γ*Δt) + γ/σ^2 - θ)
    for k in 1:nΛ
        m = Λ[k]
        for n in 0:m
            idx = n+1
            wms_prime[idx] += wms[k]*pdf(Binomial(m, p), n)
        end
    end
    return wms_prime
end

function d_CIR(m::Integer, n::Integer)
    return m .+ n
end
function e_CIR(θ1::Real, θ2::Real, β::Real)
    return θ1 + θ2 - β
end

max_stars_repo_name: konkam/DualOptimalFiltering

max_issues_repo_name: konkam/DualOptimalFiltering

max_forks_repo_name: konkam/DualOptimalFiltering


text:
function ekepler(m,ecc)
eps=1e-12
#real*8 e,e0,eps,m,ms,pi2,f0,f1,f2,f3,d1,d2,d3
# This routine solves Kepler's equation for E as a function of (e,M)
# using the procedure outlined in Murray & Dermott:
pi2=2.0*pi
ms=mod(m,pi2)
d3=1e10
e0=ms+ecc*0.85*sin(ms)/abs(sin(ms))
while abs(d3) > eps
  f3=ecc*cos(e0)
  f2=ecc*sin(e0)
  f1=1.0-f3
  f0=e0-ms-f2
  d1=-f0/f1
  d2=-f0/(f1+0.5*d1*f2)
  d3=-f0/(f1+d2*0.5*(f2+d2*f3/3.))
  e0=e0+d3
end
ekep=e0+m-ms
return ekep
end

function kepler(m,ecc)
@assert(ecc >= 0.0)
@assert(ecc <= 1.0)
f=m
if ecc > 0
  ekep=ekepler(m,ecc)
#  println(m-ekep+ecc*sin(ekep))
  f=2.0*atan(sqrt((1.0+ecc)/(1.0-ecc))*tan(0.5*ekep))
end
return f
end


text:
using Base.Test

using TMCM3110

TMCM3110.get_portname_of_TMCM3110()
#smc = LibSerialPort.open_serial_port(TMCM3110.TMCM3110_PORTNAME, TMCM3110.BAUTRATE)

if (TMCM3110.TMCM3110_PORTNAME == "")
  info("TMCM-3110 not found.")
else
  info("TMCM-3110 found. Portname: $(TMCM3110.TMCM3110_PORTNAME)")
  smc = LibSerialPort.open_serial_port(TMCM3110.TMCM3110_PORTNAME, TMCM3110.BAUTRATE)
  TMCM3110.list_all_axis_parameters(smc)
end


text:
#ДАНО: Робот - в юго-западном углу поля, на котором расставлено некоторое количество маркеров
#РЕЗУЛЬТАТ: Функция вернула значение средней температуры всех замаркированных клеток
function srednee_t(r::Robot)
    score = 0
    count = 0
    side = 3
    while true
        while !isborder(r, HorizonSide(side))
            if ismarker(r)
                score+=temperature(r)
                count+=1
            end
            move!(r, HorizonSide(side))
        end
        if ismarker(r)
            score+=temperature(r)
            count+=1
        end
        if isborder(r, HorizonSide(2))
            break
        end
        move!(r, HorizonSide(2))
        #println(side)
        side = (side+2) % 4
    end
    println(score/count, " ", count)
end

max_stars_repo_name: listova-nastya/mirea-progs

max_issues_repo_name: listova-nastya/mirea-progs

max_forks_repo_name: listova-nastya/mirea-progs


text:
#contains functions related to building, loading, and saving metadata
#currently metadata is a vector of MetaDictionary, but it could eventually
#become its own object

#reads in the metadata and creates type tuples
function readMetaIndex(;indexPath::String = DOCUMENTATION_PATH,
  fileName::String = META_INDEX_NAME)::DataFrame

  df::DataFrame = CSV.read("$indexPath\\$fileName.csv",# strings=:raw,
    categorical=false)
  #df::DataFrame = load("$indexPath\\$fileName.csv") |> DataFrame

  return df

end

#froms the vector of Metadata Dictionaries
function formMetadata(indexPath::String = DOCUMENTATION_PATH,
  indexName::String = META_INDEX_NAME)::Vector{MetaDictionary}

  #first get a list of the paths
  indexDF::DataFrame = readMetaIndex()
  NFiles::Int = size(indexDF,1)

  metadata::Vector{MetaDictionary} = (i::Int -> MetaDictionary(indexDF[i,:])).(1:NFiles)

  return metadata
end

#checks for and removes metadata fields not actually in the file
function validateMetaDictionary!(meta::MetaDictionary, keeperList::Vector{Symbol};
  nccsPath::String = NCCS_PATH,  streamer::Function = I_CSV_STREAMER)::MetaDictionary

  #these are the column names that we expect
  expectedCols::Vector{Symbol} = collect(keys(meta.dictionary))

  path::String = "$nccsPath\\$(meta.fileName).csv.gz"
  if isfile(path)
    #try
      #iStream = streamer(path)
      b::IOBuffer = open(GzipDecompressorStream, path) do f
        b = IOBuffer("$(readline(f))\n$(readline(f))")
        (b)
      end
      df::DataFrame = CSV.read(b,# strings=:raw,
        categorical=false, silencewarnings=true, strict=false)
      #df::DataFrame = load(iStream, nrows=2) |> DataFrame
      #close(iStream)

      #println("names: $(names(df))")

      #check first that the column names are correct
      actualCols::Vector{Symbol} = (s::Symbol->Symbol(lowercase(string(s)))).(names(df))

      #get the list of files we don't want
      expectedNotFound::Vector{Symbol} = setdiff(expectedCols, actualCols)
      expectedNotWanted::Vector{Symbol} = setdiff(expectedCols, keeperList)
      fieldsToRemove::Vector{Symbol} = unique([expectedNotFound; expectedNotWanted])

      #remove the offending fields
      if length(fieldsToRemove) > 0

        removeFields!(meta, fieldsToRemove)
        println("Validated $(meta.fileName).
          Removed $(length(fieldsToRemove)) fields from metadata.")
      else
        println("Validated $(meta.fileName).")
      end

  else
    @warn "path $path does not exist."
  end

  return meta

end

function validateMetadata(metadata::Vector{MetaDictionary};
    indexPath::String = DOCUMENTATION_PATH, keeperName::String=KEEPER_NAME)::Vector{MetaDictionary}

  metadata = deepcopy(metadata) #its hard to do this without some modifications AND copies
                                #hence err on the side of caution here

  metadata = removeDictionary(metadata, "nccs_core_2013_co_legacy")
  NMeta::Int = length(metadata)

  keeperDF::DataFrame = CSV.read("$indexPath\\$keeperName.csv",# strings=:raw,
    categorical=false)
  keeperList::Vector{Symbol} = (Symbol).(keeperDF[:keeper])

  #this can take a while so multi-thread it
  println("Beginning file validation")

  futures::Vector{Future} = Vector{Future}()

  #use this to check orders
  metaNames::Vector{String} = (m::MetaDictionary->m.fileName).(metadata)

  #assign and retrieve the work
  for i::Int ∈ 1:NMeta
    push!(futures, (@spawn validateMetaDictionary!(metadata[i], keeperList)))
  end
  metadata .= (fetch).(futures)

  return metadata
end

#this function either forms or loads the metadata
function getMetadata(;refreshMetadata::Bool = false, metadataPath::String = DOCUMENTATION_PATH,
  metadataName::String = METADATA_NAME)

  metadataFilePath::String = "$metadataPath\\$metadataName.jls"

  #load the metadata if we don't need to redo the validation
  if refreshMetadata || !(isfile(metadataFilePath))
    metadata::Vector{MetaDictionary}= formMetadata()
    metadata = validateMetadata(metadata)

    oStream::IOStream = open(metadataFilePath, "w")
    serialize(oStream, metadata)
    close(oStream)
  else
    iStream::IOStream = open(metadataFilePath)
    metadata = deserialize(iStream)
    close(iStream)
  end

  return metadata
end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
struct Fourier{T} <: Basis{T} end
Fourier() = Fourier{Float64}()

==(::Fourier, ::Fourier) = true

axes(F::Fourier) = (Inclusion(ℝ), _BlockedUnitRange(1:2:∞))

function getindex(F::Fourier{T}, x::Real, j::Int)::T where T
    isodd(j) && return cos((j÷2)*x)
    sin((j÷2)*x)
end

### transform
checkpoints(::Fourier{T}) where T = T[1.223972,3.14,5.83273484]

fouriergrid(T, n) = convert(T,π)*collect(0:2:2n-2)/n

function grid(Pn::SubQuasiArray{T,2,<:Fourier,<:Tuple{<:Inclusion,<:AbstractUnitRange}}) where T
    kr,jr = parentindices(Pn)
    n = maximum(jr)
    fouriergrid(T, n)
end

"""
Gives a shuffled version of the real FFT, with order
1,sin(θ),cos(θ),sin(2θ)…
"""
struct ShuffledRFFT{T,Pl<:Plan} <: Factorization{T}
    plan::Pl
end

size(F::ShuffledRFFT, k) = size(F.plan,k)
size(F::ShuffledRFFT) = size(F.plan)

ShuffledRFFT{T}(p::Pl) where {T,Pl<:Plan} = ShuffledRFFT{T,Pl}(p)
ShuffledRFFT{T}(n, d...) where T = ShuffledRFFT{T}(FFTW.plan_r2r(Array{T}(undef, n), FFTW.R2HC, d...))

function _shuffledrfft_postscale!(_, ret::AbstractVector{T}) where T
    n = length(ret)
    lmul!(convert(T,2)/n, ret)
    ret[1] /= 2
    iseven(n) && (ret[n÷2+1] /= 2)
    negateeven!(reverseeven!(interlace!(ret,1)))
end

function _shuffledrfft_postscale!(d::Number, ret::AbstractMatrix{T}) where T
    if isone(d)
        n = size(ret,1)
        lmul!(convert(T,2)/n, ret)
        ldiv!(2, view(ret,1,:))
        iseven(n) && ldiv!(2, view(ret,n÷2+1,:))
        for j in axes(ret,2)
            negateeven!(reverseeven!(interlace!(view(ret,:,j),1)))
        end
    else
        n = size(ret,2)
        lmul!(convert(T,2)/n, ret)
        ldiv!(2, view(ret,:,1))
        iseven(n) && ldiv!(2, view(ret,:,n÷2+1))
        for k in axes(ret,1)
            negateeven!(reverseeven!(interlace!(view(ret,k,:),1)))
        end
    end
    ret
end

function mul!(ret::AbstractArray{T}, F::ShuffledRFFT{T}, b::AbstractArray) where T
    mul!(ret, F.plan, convert(Array{T}, b))
    _shuffledrfft_postscale!(F.plan.region, ret)
end

*(F::ShuffledRFFT{T}, b::AbstractVecOrMat) where T = mul!(similar(b, T), F, b)

factorize(L::SubQuasiArray{T,2,<:Fourier,<:Tuple{<:Inclusion,<:OneTo}}) where T =
    TransformFactorization(grid(L), ShuffledRFFT{T}(size(L,2)))
factorize(L::SubQuasiArray{T,2,<:Fourier,<:Tuple{<:Inclusion,<:OneTo}}, d) where T =
    TransformFactorization(grid(L), ShuffledRFFT{T}((size(L,2),d),1))

factorize(L::SubQuasiArray{T,2,<:Fourier,<:Tuple{<:Inclusion,<:BlockSlice}},d...) where T =
    ProjectionFactorization(factorize(parent(L)[:,OneTo(size(L,2))],d...),parentindices(L)[2])

import BlockBandedMatrices: _BlockSkylineMatrix

@simplify function *(A::QuasiAdjoint{<:Any,<:Fourier}, B::Fourier)
    TV = promote_type(eltype(A),eltype(B))
    PseudoBlockArray(Diagonal(Vcat(2convert(TV,π),Fill(convert(TV,π),∞))), (axes(A,1),axes(B,2)))
end

@simplify function *(D::Derivative, F::Fourier)
    TV = promote_type(eltype(D),eltype(F))
    Fourier{TV}()*_BlockArray(Diagonal(Vcat([reshape([0.0],1,1)], (1.0:∞) .* Fill([0 -one(TV); one(TV) 0], ∞))), (axes(F,2),axes(F,2)))
end

function broadcasted(::LazyQuasiArrayStyle{2}, ::typeof(*), c::BroadcastQuasiVector{<:Any,typeof(cos),<:Tuple{<:Inclusion{<:Any,RealNumbers}}}, F::Fourier)
    axes(c,1) == axes(F,1) || throw(DimensionMismatch())
    T = promote_type(eltype(c), eltype(F))
    # Use LinearAlgebra.Tridiagonal for now since broadcasting support not complete for LazyBandedMatrices.Tridiagonal
    F*mortar(LinearAlgebra.Tridiagonal(Vcat([reshape([0; one(T)],2,1)], Fill(Matrix(one(T)/2*I,2,2),∞)),
                        Vcat([zeros(T,1,1)], Fill(Matrix(zero(T)I,2,2),∞)),
                        Vcat([[0 one(T)/2]], Fill(Matrix(one(T)/2*I,2,2),∞))))
end

function broadcasted(::LazyQuasiArrayStyle{2}, ::typeof(*), s::BroadcastQuasiVector{<:Any,typeof(sin),<:Tuple{<:Inclusion{<:Any,RealNumbers}}}, F::Fourier)
    axes(s,1) == axes(F,1) || throw(DimensionMismatch())
    T = promote_type(eltype(s), eltype(F))
    # Use LinearAlgebra.Tridiagonal for now since broadcasting support not complete for LazyBandedMatrices.Tridiagonal
    F*mortar(LinearAlgebra.Tridiagonal(Vcat([reshape([one(T); 0],2,1)], Fill([0 one(T)/2; -one(T)/2 0],∞)),
                        Vcat([zeros(T,1,1)], Fill(Matrix(zero(T)*I,2,2),∞)),
                        Vcat([[one(T)/2 0]], Fill([0 -one(T)/2; one(T)/2 0],∞))))
end

# support cos.(θ) .* F
Base.broadcasted(::LazyQuasiArrayStyle{2}, ::typeof(*), c::Broadcasted{<:Any,<:Any,typeof(cos),<:Tuple{<:Inclusion{<:Any,RealNumbers}}}, F::Fourier) = materialize(c) .* F
Base.broadcasted(::LazyQuasiArrayStyle{2}, ::typeof(*), c::Broadcasted{<:Any,<:Any,typeof(sin),<:Tuple{<:Inclusion{<:Any,RealNumbers}}}, F::Fourier) = materialize(c) .* F


text:
using POMDPs, Crux, Flux, POMDPGym, Random, Distributions, POMDPPolicies
Crux.set_function("isfailure", POMDPGym.isfailure)

# Define the disturbance distribution based on a normal distribution
px = MvNormal2([0f0], [0.5f0])

# Construct the MDP
mdp = AdditiveAdversarialMDP(ContinuousBanditMDP(), px)
S = state_space(mdp)
N = 10000

# construct the model
QSA() = ContinuousNetwork(Chain(Dense(2, 64, relu), Dense(64, 64, relu), Dense(64, 1)))
Pf() = ContinuousNetwork(Chain(Dense(2, 64, relu), Dense(64, 64, relu), Dense(64, 1, sigmoid)))
A() = ContinuousNetwork(Chain(Dense(1, 64, relu), Dense(64, 64, relu), Dense(64, 1, tanh), x -> 2f0 * x), 1)

G() = GaussianPolicy(A(), zeros(Float32, 1))

Protag() = ActorCritic(A(), DoubleNetwork(QSA(), QSA()))
Antag() = ActorCritic(G(), Pf())

AdvPol(p = Protag()) = AdversarialPolicy(p, Antag())

𝒮_td3 = TD3(;π=ActorCritic(A(), DoubleNetwork(QSA(), QSA())), S=S, N=50000, buffer_size=Int(1e5), buffer_init=1000)
π_td3 = solve(𝒮_td3, mdp)

# solve with IS
𝒮_isarl = ISARL_Continuous(π=AdvPol(),
                           S=S,
                           N=50000,
                           π_explore=GaussianNoiseExplorationPolicy(Crux.LinearDecaySchedule(0., 0.0, floor(Int, 15000))),
                           px=px,
                           return_at_episode_end=false,
                           buffer_size=Int(1e5),
                           buffer_init=1000,)

π_isarl = solve(𝒮_isarl, mdp)

plot(-3:0.1:3, [value(antagonist(𝒮_isarl.π), [0], [i])[1] for i=-3:0.1:3], label="Probability of failure")

plot!(-3:0.1:3, [exp(logpdf(antagonist(𝒮_isarl.π).A, [0], [i])[1]) for i=-3:0.1:3], label="Failure Policy")

plot!(-3:0.1:3, [exp.(logpdf(px, i))[1] for i=-3:0.1:3], label="Px")

logpdf(px, 0.1)


text:
function collatz(n, count=1)
  if n == 1
    return count
  end
  if iseven(n)
    n = Int(n / 2)
  else
    n = 3n + 1
  end
  count += 1
  collatz(n, count)
end

maxc = 0
maxnum = 0
for i = 1:1000_000
  global maxc, maxnum
  newc = collatz(i)
  if newc > maxc
    maxc = newc
    maxnum = i
  end
end
println(maxnum)

max_stars_repo_name: sujimodern/project-euler-jl

max_issues_repo_name: sujimodern/project-euler-jl

max_forks_repo_name: sujimodern/project-euler-jl


text:
function sort_sections(chn::Chains)
    smap = keys(chn.name_map)
    section_list = Vector{Symbol}(undef, length(smap))
    indx = 1
    if :parameters in smap
        section_list[1] = :parameters
        indx += 1
    end
    if :internals in smap
        section_list[end] = :internals
    end
    for par in smap
        if !(par in [:parameters, :internals])
            section_list[indx] = par
            indx += 1
        end
    end
    return section_list
end

"""

# Array

Array constructor from a Chains object. Returns 3 dimensionsal
array or an Array of 2 dimensional Arrays. If only a single parameter is selected for
inclusion, a dimension is dropped in both cases, as is e.g. required by cde(), etc.

### Method
```julia
  Array(
    chn::Chains,
    sections::Vector{Symbol};
    append_chains::Bool,
    remove_missing_union::Bool
  )
```

### Required arguments
```julia
* `chn` : Chains object to convert to an Array
```

### Optional arguments
```julia
* `sections = Symbol[]` : Sections from the Chains object to be included
* `append_chains = true`  : Append chains into a single column
* `remove_missing_union = true`  : Convert Union{Missing, Real} to Float64
```

### Examples
```julia
* `Array(chns)` : Array with chain values are appended
* `Array(chns[:par])` : Array with selected parameter chain values are appended
* `Array(chns, [:parameters])`  : Array with only :parameter section
* `Array(chns, [:parameters, :internals])`  : Array includes multiple sections
* `Array(chns, append_chains=false)`  : Array of Arrays, each chain in its own array
* `Array(chns, remove_missing_union=false)` : No conversion to remove missing values
```

"""
function Base.Array(chn::Chains,
        sections::Union{Symbol, Vector{Symbol}}=Symbol[:parameters];
        append_chains=true,
        showall=false,
        sorted=false
    )
    sections = showall ? keys(chn.name_map) : sections
    chn = Chains(chn, sections)

    arr = if append_chains
        mapreduce(i -> chn.value.data[:,:,i], vcat, 1:size(chn, 3))
    else
        map(i -> chn.value.data[:,:,i], 1:size(chn, 3))
    end

    arr = if append_chains
        reshape(arr, (size(arr, 1), size(arr, 2)))
    else
        arr
    end

    if size(arr, 2) == 1
        return map(identity, arr[:,1])
    else
        return map(identity, arr)
    end
end


text:
##### update coordinates of each individual using RK4 integration
function plankton_advection!(plank, velos, g::AbstractGrid, vel₀, vel½, vel₁, ΔT, arch::Architecture)
    vel_interpolate!(velos.u1, velos.v1, velos.w1, plank.x, plank.y, plank.z, plank.ac,
                     vel₀.u.data, vel₀.v.data, vel₀.w.data, g, arch)

    ##### add up intermediate velocities
    velos.u2 .= velos.u1
    velos.v2 .= velos.v1
    velos.w2 .= velos.w1

    calc_coord!(velos, plank, velos.u1, velos.v1, velos.w1, plank.ac, ΔT, 0.5, arch)
    particle_boundaries!(velos, plank.ac, g, arch)

    ##### stage 2
    vel_interpolate!(velos.u1, velos.v1, velos.w1, velos.x, velos.y, velos.z, plank.ac,
                     vel½.u.data, vel½.v.data, vel½.w.data, g, arch)

    ##### add up intermediate velocities
    velos.u2 .+= velos.u1 .* 2
    velos.v2 .+= velos.v1 .* 2
    velos.w2 .+= velos.w1 .* 2

    calc_coord!(velos, plank, velos.u1, velos.v1, velos.w1, plank.ac, ΔT, 0.5, arch)
    particle_boundaries!(velos, plank.ac, g, arch)

    ##### stage 3
    vel_interpolate!(velos.u1, velos.v1, velos.w1, velos.x, velos.y, velos.z, plank.ac,
                     vel½.u.data, vel½.v.data, vel½.w.data, g, arch)

    ##### add up intermediate velocities
    velos.u2 .+= velos.u1 .* 2
    velos.v2 .+= velos.v1 .* 2
    velos.w2 .+= velos.w1 .* 2

    calc_coord!(velos, plank, velos.u1, velos.v1, velos.w1, plank.ac, ΔT, 1.0, arch)
    particle_boundaries!(velos, plank.ac, g, arch)

    ##### stage 4
    vel_interpolate!(velos.u1, velos.v1, velos.w1, velos.x, velos.y, velos.z, plank.ac,
                     vel₁.u.data, vel₁.v.data, vel₁.w.data, g, arch)

    ##### add up intermediate velocities
    velos.u2 .+= velos.u1
    velos.v2 .+= velos.v1
    velos.w2 .+= velos.w1

    ##### calculate final velocities
    velos.u2 .= velos.u2 ./ 6
    velos.v2 .= velos.v2 ./ 6
    velos.w2 .= velos.w2 ./ 6

    calc_coord!(plank, plank, velos.u2, velos.v2, velos.w2, plank.ac, ΔT, 1.0, arch)
    particle_boundaries!(plank, plank.ac, g, arch)

    return nothing
end

# ##### update coordinates of each individual using quasiAB2 integration
# function plankton_advection!(plank, velos, g::AbstractGrid, χ, vel₁, ΔT, arch::Architecture)
#     ##### time step n
#     vel_interpolate!(velos.u1, velos.v1, velos.w1, plank.x, plank.y, plank.z, plank.ac,
#                      vel₁.u.data, vel₁.v.data, vel₁.w.data, g, arch)

#     ##### AB2
#     calc_vel_ab2!(velos, χ, arch)
#     calc_coord!(plank, plank, velos.u1, velos.v1, velos.w1, plank.ac, ΔT, 1.0, arch)

#     ##### time step n-1
#     velos.u2 .= velos.u1
#     velos.v2 .= velos.v1
#     velos.w2 .= velos.w1

#     return nothing
# end

# ##### update coordinates of each individual using Explicit Euler (aka Euler Forward) integration
# function plankton_advection!(plank, velos, g::AbstractGrid, vel, ΔT, arch::Architecture)
#     vel_interpolate!(velos.u1, velos.v1, velos.w1, plank.x, plank.y, plank.z, plank.ac,
#                      vel.u.data, vel.v.data, vel.w.data, g, arch)

#     calc_coord!(plank, plank, velos.u1, velos.v1, velos.w1, plank.ac, ΔT, 1.0, arch)
#     particle_boundaries!(plank, plank.ac, g, arch)

#     return nothing
# end


text:
using Quantica: nvertices, nedges, nsimplices

@testset "meshes" begin
    m = cuboid((0,1), (0,2), subticks = (10, 20))
    @test nvertices(m) == 200
    @test_throws MethodError cuboid((SA[1,2], SA[2,3]))
end

max_stars_repo_name: pablosanjose/Quantica

max_issues_repo_name: pablosanjose/Quantica

max_forks_repo_name: pablosanjose/Quantica


text:
#=
  Coins puzzle in Julia + ConstraintSolver.jl.

  Problem from
  Tony Hürlimann: "A coin puzzle - SVOR-contest 2007"
  http://www.svor.ch/competitions/competition2007/AsroContestSolution.pdf
  """
  In a quadratic grid (or a larger chessboard) with 31x31 cells, one
  should place coins in such a way that the following conditions are
  fulfilled:
    1. In each row exactly 14 coins must be placed.
    2. In each column exactly 14 coins must be placed.
    3. The sum of the quadratic horizontal distance from the main
       diagonal of all cells containing a coin must be as small as possible.
    4. In each cell at most one coin can be placed.

   The description says to place 14x31 = 434 coins on the chessboard
   each row containing 14 coins and each column also containing 14 coins.
  """

  Note: This problem is quite hard for CP solvers. A MIP solver solves
  the (14,31) problem in millis.

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

using ConstraintSolver, JuMP
using Cbc
using GLPK
const CS = ConstraintSolver
include("constraints_utils.jl")

function coins_grid(n=31,c=14)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer) # , "logLevel" => 0)
    model = Model(optimizer_with_attributes(CS.Optimizer,   # "all_solutions"=> true,
                                                            # "all_optimal_solutions"=>true,
                                                            "logging"=>[],

                                                            # "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # "simplify"=>false,
                                                            "simplify"=>true,

                                                            "time_limit"=>110,
                                                            # "lp_optimizer" => cbc_optimizer, # 10.7s
                                                            "lp_optimizer" => glpk_optimizer, # 1.1s
                                        ))

    @variable(model, x[1:n,1:n], Bin)
    @variable(model, 0 <= s <= n*n*n, Int)

    @constraint(model,s == sum([(x[i,j] * abs(i-j)*abs(i-j)) for i in 1:n, j in 1:n]))
    for i in 1:n
        @constraint(model, c == sum(x[i,:]))
        @constraint(model, c == sum(x[:,i]))
    end

    @objective(model, Min, s)

    # Solve the problem
    println("solve")
    optimize!(model)

    status = JuMP.termination_status(model)
    println("status:$status")
    if status == MOI.OPTIMAL
        x = convert.(Integer,JuMP.value.(x))
        s = convert.(Integer,JuMP.value.(s))
        for i in 1:n
            println(x[i,:])
        end
        println("s:$s")

        num_sols = MOI.get(model, MOI.ResultCount())
        println("\nnum_sols:$num_sols\n")
    end
end

n = 31
c = 14
coins_grid(n,c)

max_stars_repo_name: Wikunia/hakank

max_issues_repo_name: Wikunia/hakank

max_forks_repo_name: Wikunia/hakank


text:
nt = 1000
seed = 99
m = ACQR_simula(nt,seed)
y = m["y"]

############################### starting point
nx = 4
i = nx+1
Ai,Ci,Qi,Ri,Si = ACQRS_sub(y,nx,i)
m1i = Ci\y[:,1:1]
P1i = 0*Ai

#######################
max_iter=10
tol=1e-6
txo=false
print("Testing ACQR_em: ")
A,C,Q,R,m1,P1,loglik,aic = ACQR_em(y,Ai,Ci,Qi,Ri,m1i,P1i,max_iter,tol,txo)
if sum(A.^2) > 1e-6
    println("OK")
end


text:
@testset "GeodesicOnTheSphere" begin
    @polyvar x y z
    f = x + y + 3z
    g = x + 2z
    h = y + 2z

    @test GeodesicOnTheSphere(f, g) isa GeodesicOnTheSphere{Float64}
    @test GeodesicOnTheSphere([f, f], [g, h]) isa GeodesicOnTheSphere{Float64}
    @test GeodesicOnTheSphere{Float64}([f, g], [f, h]) isa GeodesicOnTheSphere{Float64}
    @test GeodesicOnTheSphere{Float64}(f, g) isa GeodesicOnTheSphere{Float64}
    @test GeodesicOnTheSphere([FP.Polynomial(f)], [FP.Polynomial{Complex128}(f)]) isa GeodesicOnTheSphere{Complex128}
    @test_throws AssertionError GeodesicOnTheSphere{Complex128}([f], [g])
    @test_throws AssertionError GeodesicOnTheSphere{Complex128}([f], [f, f])

    H = GeodesicOnTheSphere{Float64}([f, g], [f, h])
    K = GeodesicOnTheSphere([f, g], [f, h])
    @test promote_type(typeof(H), typeof(K)) == GeodesicOnTheSphere{Float64}
    @test convert(GeodesicOnTheSphere{Float64}, K) isa GeodesicOnTheSphere{Float64}

    w = [1.0, 2.0, 2.0]
    @test evaluate(H, w, 1.0) == [9/4, 5/4]
    @test H(w, 1.0) == [9/4, 5/4]
    u = zeros(2)
    evaluate!(u, H, w, 1.0)
    @test u == [9/4, 5/4]

    cfg = PolynomialHomotopyConfig(H)
    @test evaluate(H, w, 1.0, cfg) == [9/4, 5/4]
    u = zeros(2)
    evaluate!(u, H, w, 1.0, cfg)
    @test u == [9/4, 5/4]

    @test jacobian(H, rand(3), 0.0, cfg) == [1 1 3; 0 1 2]./4
    @test jacobian(H, rand(3), 1.0, cfg) == [1 1 3; 1 0 2]./4
    u = zeros(2, 3)
    @test jacobian!(u, H, rand(3), 0.0, cfg) == [1 1 3; 0 1 2]./4
    @test u == [1 1 3; 0 1 2]./4

    r = JacobianDiffResult(cfg)
    w = rand(3)
    jacobian!(r, H, w, 0.24, cfg)
    @test value(r) == H(w, 0.24)
    @test jacobian(r) == jacobian(H, w, 0.24, cfg)

    # TODO: This test is bonkers atm
    @test dt(H, [1, 2, 2.0], 0.0, cfg) != [0, 0]
    u = zeros(2)
    dt!(u, H, [1, 2, 2.0], 0.0, cfg)
    @test u == dt(H, [1, 2, 2.0], 0.0, cfg)
    @test string(H) == "Homotopies.GeodesicOnTheSphere{Float64} with 2 polynomials in 3 variables.\n"

    r = DtDiffResult(cfg)
    w = rand(3)
    dt!(r, H, w, 0.24, cfg)
    @test value(r) == H(w, 0.24)
    @test dt(r) == dt(H, w, 0.24, cfg)

    @test promote_type(GeodesicOnTheSphere{Float64}, Complex128) == GeodesicOnTheSphere{Complex128}
    @test promote_type(GeodesicOnTheSphere{Float64}, GeodesicOnTheSphere{Complex128}) == GeodesicOnTheSphere{Complex128}

    N=weylnorm(H)
    @test N(0.0)==1.0
    @test N(0.5)==1.0

    H = GeodesicOnTheSphere(x^2+y+z, z^2+2+x+y^2)
    @test H == H

    HH = homogenize(H)
    @test ishomogenous(H) == true
    @test ishomogenized(H) == true
    @test ishomogenous(HH)
    @test ishomogenized(HH)

    @test length(H) == 1
    @test nvariables(H) == 4
end


text:
function ∂ end

function ∂(::GridType{:Primal},grid::Grid{Dim},dir::Direction{D}, F::AbstractArray) where {Dim,D}
    shifts = -1 .* (ntuple(identity,Val(Dim)) .== D)
	if D > Dim return zeros(CFloat, size(grid)) end
    (circshift(F,shifts) - F) / spacing(grid,dir)
end

function ∂(::GridType{:Dual},grid::Grid{Dim},dir::Direction{D}, F::AbstractArray) where {Dim,D}
    shifts = 1 .* (ntuple(identity,Val(Dim)) .== D)
	 if D > Dim return zeros(CFloat, size(grid)) end
    (circshift(F,shifts) - F) / spacing(grid,dir)
end


text:
@testset "prob1" begin
    @markov midplane(lam, mu) begin
        @tr begin
            :up => :down, lam
            :down => :up, mu
        end
        @reward :r begin
            :up, 1
        end
    end

    begin
        lam = 0.4
        mu = 1.5
        delta = 1.5
    end

    m = ctmc(midplane(lam, mu))
    println(sprob(m))
end

@testset "sprob2" begin
    @markov midplane(lam, mu) begin
        @tr begin
            :up => :down, lam
            :down => :up, mu
        end
        @reward :r begin
            :up, 1
        end
    end

    begin
        lam = 0.4
        mu = 1.5
        delta = 1.5
    end

    m = ctmc(midplane(lam, mu), :DenseCTMC)
    println(sprob(m))
end

@testset "sprob2" begin
    @markov midplane(lam, mu) begin
        @tr begin
            :up => :down, lam
            :down => :up, mu
        end
        @reward :r begin
            :up, 1
        end
    end

    begin
        lam = 0.4
        mu = 1.5
        delta = 1.5
    end

    m = ctmc(midplane(lam, mu), :DenseCTMC)
    println(sprob(m, states=[:up]))
end

@testset "sprob4" begin
    @markov midplane(lam, mu) begin
        @tr begin
            :up => :down, lam
            :down => :up, mu
        end
        @reward :r begin
            :up, 1
        end
    end

    @bind begin
        lam = 0.4
        mu = 1.5
        delta = 1.5
    end

    m = ctmc(midplane(lam, mu), :DenseCTMC)
    println(seval(sprob(m)))
    println(seval(sprob(m), :lam))
    println(seval(sprob(m), (:lam, :lam)))
end

@testset "sprob5" begin
    @markov midplane(lam, mu) begin
        @tr begin
            :up => :down, lam
            :down => :up, mu
        end
        @reward :r begin
            :up, 1
        end
    end

    @bind begin
        lam = 0.4
        mu = 1.5
        delta = 1.5
    end

    m = ctmc(midplane(lam, mu), :DenseCTMC)
    println(seval(sprob(m, states=[:up])))
    println(seval(sprob(m, states=[:up]), :lam))
    println(seval(sprob(m, states=[:up]), (:lam, :lam)))
end


text:
using Test
import MathOptInterface
const MOI = MathOptInterface

include("../model.jl")

@testset "Scalar" begin
    model = Model{Float64}()
    x = MOI.add_variable(model)
    @testset "SingleVariable" begin
        f = MOI.SingleVariable(x)
        ci = MOIU.add_scalar_constraint(model, f, MOI.EqualTo(1.0),
                                        allow_modify_function = false)
        @test MOI.get(model, MOI.ConstraintFunction(), ci) == f
        @test MOI.get(model, MOI.ConstraintSet(), ci) == MOI.EqualTo(1.0)
    end
    @testset "ScalarAffineFunction" begin
        f = MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x)], 2.0)
        g = MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x)], 0.0)
        ci = MOIU.add_scalar_constraint(model, f, MOI.EqualTo(3.0))
        @test f.constant == 2.0
        @test MOI.get(model, MOI.ConstraintFunction(), ci) ≈ g
        @test MOI.get(model, MOI.ConstraintSet(), ci) == MOI.EqualTo(1.0)
        ci = MOIU.add_scalar_constraint(model, f, MOI.Interval(-1.0, 1.0),
                                        allow_modify_function = true)
        @test f.constant == 0.0
        @test MOI.get(model, MOI.ConstraintFunction(), ci) ≈ g
        @test MOI.get(model, MOI.ConstraintSet(), ci) == MOI.Interval(-3.0,
                                                                      -1.0)
    end
end


text:
module SamplingTaskFarm

using IntervalSets
using DelimitedFiles

using Sockets
using Dates
using PrettyTables
using IOUtils

hostname() = strip(String(read(`hostname`)))
const localhost = ip"127.0.0.1"

include("samplers.jl")
include("task_farm.jl")

export ..

end


text:
#############################################
# Helpers
#############################################
function exchange(a,b)
	 temp = a
         a=b
         b=temp
	 return (a,b)
end

function sparseKeepZero(I::AbstractVector{Ti},
    J::AbstractVector{Ti},
    V::AbstractVector{Tv},
    nrow::Integer, ncol::Integer) where {Tv,Ti<:Integer}
    N = length(I)
    if N != length(J) || N != length(V)
        throw(ArgumentError("triplet I,J,V vectors must be the same length"))
    end
    if N == 0
        return spzeros(eltype(V), Ti, nrow, ncol)
    end

    # Work array
    Wj = Array{Ti}(undef,max(nrow,ncol)+1)
    # Allocate sparse matrix data structure
    # Count entries in each row
    Rnz = zeros(Ti, nrow+1)
    Rnz[1] = 1
    nz = 0
    for k=1:N
        iind = I[k]
        iind > 0 || throw(ArgumentError("all I index values must be > 0"))
        iind <= nrow || throw(ArgumentError("all I index values must be ≤ the number of rows"))
        Rnz[iind+1] += 1
        nz += 1
    end
    Rp = cumsum(Rnz)
    Ri = Array{Ti}(undef,nz)
    Rx = Array{Tv}(undef,nz)

    # Construct row form
    # place triplet (i,j,x) in column i of R
    # Use work array for temporary row pointers
    @simd for i=1:nrow; @inbounds Wj[i] = Rp[i]; end
    @inbounds for k=1:N
        iind = I[k]
        jind = J[k]
        jind > 0 || throw(ArgumentError("all J index values must be > 0"))
        jind <= ncol || throw(ArgumentError("all J index values must be ≤ the number of columns"))
        p = Wj[iind]
        Vk = V[k]
        Wj[iind] += 1
        Rx[p] = Vk
        Ri[p] = jind
    end

    # Reset work array for use in counting duplicates
    @simd for j=1:ncol; @inbounds Wj[j] = 0; end

    # Sum up duplicates and squeeze
    anz = 0
    @inbounds for i=1:nrow
        p1 = Rp[i]
        p2 = Rp[i+1] - 1
        pdest = p1
        for p = p1:p2
            j = Ri[p]
            pj = Wj[j]
            if pj >= p1
                Rx[pj] = Rx[pj] + Rx[p]
            else
                Wj[j] = pdest
                if pdest != p
                    Ri[pdest] = j
                    Rx[pdest] = Rx[p]
                end
                pdest += one(Ti)
            end
        end
        Rnz[i] = pdest - p1
        anz += (pdest - p1)
    end

    # Transpose from row format to get the CSC format
    RiT = Array{Ti}(undef,anz)
    RxT = Array{Tv}(undef,anz)

    # Reset work array to build the final colptr
    Wj[1] = 1
    @simd for i=2:(ncol+1); @inbounds Wj[i] = 0; end
    @inbounds for j = 1:nrow
        p1 = Rp[j]
        p2 = p1 + Rnz[j] - 1
        for p = p1:p2
            Wj[Ri[p]+1] += 1
        end
    end
    RpT = cumsum(Wj[1:(ncol+1)])

    # Transpose
    @simd for i=1:length(RpT); @inbounds Wj[i] = RpT[i]; end
    @inbounds for j = 1:nrow
        p1 = Rp[j]
        p2 = p1 + Rnz[j] - 1
        for p = p1:p2
            ind = Ri[p]
            q = Wj[ind]
            Wj[ind] += 1
            RiT[q] = j
            RxT[q] = Rx[p]
        end
    end

    return SparseMatrixCSC(nrow, ncol, RpT, RiT, RxT)
end

#Convert Julia indices to C indices
function convert_to_c_idx(indicies)
    for i in 1:length(indicies)
        indicies[i] = indicies[i] - 1
    end
end

#Grab the subnode in hierarchical linkconstraint link
function _get_subnode(first_stage,link)
	linked_nodes = getnodes(link)
	@assert length(linked_nodes) == 2
	sub_node = setdiff(getnodes(link),[first_stage])[1]
	return sub_node
end

#categorize link constraints into hierarchical and distributed
function _identify_linkconstraints(graph::OptiGraph)
    links_connect_eq = LinkConstraint[]
    links_connect_ineq = LinkConstraint[]
    links_eq = LinkConstraint[]
    links_ineq = LinkConstraint[]

    first_stage_nodes = getnodes(graph)

    sub_nodes = setdiff(all_nodes(graph),first_stage_nodes)

    for link in all_linkconstraints(graph)
        link_constraint = constraint_object(link)
        link_nodes = getnodes(link_constraint)

        #if first_stage_node in link_nodes
		if isempty(setdiff(first_stage_nodes,link_nodes))
            @assert length(link_nodes) == 2
            if isa(link_constraint.set,MOI.EqualTo)
                push!(links_connect_eq,link_constraint)
            else
                push!(links_connect_ineq,link_constraint)
            end
        else
            if isa(link_constraint.set,MOI.EqualTo)
                push!(links_eq,link_constraint)
            else
                push!(links_ineq,link_constraint)
            end
        end
    end

    return links_connect_eq,links_connect_ineq,links_eq,links_ineq
end


text:
"""
    abstract type OliverProblem{T,N,M} end

Abstract type for constructing a boundary value problem from a recurrence
relation. The functions for coefficients and the right side dispatch.

`N` and `M` correspond to the total number of specified initial boundary conditions,
and the number specified on the end of the recurrence, respectively.
"""
abstract type OliverProblem{T,N,M} end

"""
    assembleP{T}(OP::OliverProblem{M,N}, a, b) where {M, N}

Generic description of the ``\\mathbf{P}`` matrix within the linear system of
Oliver's method. This in equation 3a of Oliver 1968.
"""
function assembleP(Tmat, OP::OliverProblem{T,N,M}, a, b) where {T,N,M}
    Psize = b - N - a
    𝐏 = BandedMatrix{Tmat}(undef, (Psize, Psize), (N-M,M))
    for row in 1:Psize
        for col in max(1, row-N+M):min(row + M, Psize)
            Δm = col - row
            𝐏[row, col] = OliverP(OP, M - Δm, a + row)
        end
    end
    return 𝐏
end
assembleP(OP::OliverProblem{T,N,M}, a, b) where {T,N,M} = assembleP(T,OP,a,b)

"""
    assembleρ{T}(OP::OliverProblem{M,N}, a, b) where {M, N}

Generic description of the ``\\vec{\\rho}`` vector within the linear system of
Oliver's method. This in equation 3b of Oliver 1968.
"""
function assembleρ(Tmat, OP::OliverProblem{T,N,M}, a, b, YBC) where {T,N,M}
    Psize = b - N - a
    ρ = zeros(Tmat, Psize)

    for row in 1:(N-M)
        Σ = zero(T)
        for s in 1:(N - M - row + 1)
            # every row, increase by 1 the sum UB, YBC arg, P arg
            Σ = Σ + OliverP(OP, N - s + 1, a + row) * YBC[a + s + row - 1]
        end
        ρ[row] = OliverR(OP, a + row) - Σ
    end

    rₛ, rₑ = (b - N - M + 1 - a), (b - N - a)
    for row in rₛ:rₑ
        row_from_end = rₑ - row + 1
        Σ = zero(T)
        for s in 0:(M - row_from_end)
            # every row, decrease by 1 the sum UB, YBC arg, P arg
            Σ = Σ + OliverP(OP, s, b - N - row_from_end + 1) *
                YBC[b - s - row_from_end + 1]
        end
        ρ[row] = OliverR(OP, a+row) - Σ
    end

    return ρ
end
assembleρ(OP::OliverProblem{T}, a, b, YBC) where T =
    assembleρ(T, OP, a, b, YBC)

"""
	OliverP(OP::OliverProblem, s, i)

Computes ``P_s(i)``, the coefficient ``s`` at solution index ``i``.
"""
function OliverP(OP::OliverProblem, s, i)
    throw(ErrorException("You need to implement an OliverP for this problem."))
end

"""
	OliverR(OP::OliverProblem, i)

Computes ``R(i)``, the right side of the recurrence relation.
"""
function OliverR(OP::OliverProblem, i)
    throw(ErrorException("You need to implement an OliverR for this problem."))
end


text:
function dist_matrix(G::SimpleDigraph)
  VV = vlist(G)
  n = length(VV)
  A = zeros(Int,n,n)
  for u=1:n
    d = dist(G,VV[u])
    for v=1:n
      A[u,v]=d[VV[v]]
    end
  end
  return A
end

function diam(G::SimpleDigraph)
  A = dist_matrix(G)
  if minimum(A) < 0
    return -1
  end
  return maximum(A)
end


text:
mutable struct ExplicitExponentialRungeKuttaSolver{stepsize_T} <: AbstractRungeKuttaSolver
    stepsize::stepsize_T
end

ExplicitExponentialRungeKuttaSolver(h::Real) = ExplicitExponentialRungeKuttaSolver(StepSize(h))
@doc (@doc ExplicitExponentialRungeKuttaSolver) EERK(args...; kwargs...) = ExplicitExponentialRungeKuttaSolver(args...; kwargs...)

############################################################################################
#                                         PRINTING                                         #
############################################################################################

"""
    show(io::IO, solver::ExplicitExponentialRungeKuttaSolver)

prints a full description of `solver` and its contents to a stream `io`.
"""
Base.show(io::IO, solver::ExplicitExponentialRungeKuttaSolver) = NSDEBase._show(io, solver)

"""
    summary(io::IO, solver::ExplicitExponentialRungeKuttaSolver)

prints a brief description of `solver` to a stream `io`.
"""
Base.summary(io::IO, solver::ExplicitExponentialRungeKuttaSolver) = NSDEBase._summary(io, solver)

############################################################################################
#                                          METHODS                                         #
############################################################################################

function (solver::ExplicitExponentialRungeKuttaSolver)(problem::AbstractInitialValueProblem)
    solve(problem, solver)
end


text:
# dimensions
nx = 2#12
nu = 1#4
nw = 0

# variables
@variables y[1:nx], x[1:nx], u[1:nu], w[1:nw]

y1 = ones(nx)
x1 = ones(nx)
u1 = ones(nu)
w1 = ones(nw)

xu1 = [x1; u1]
xuy1 = [x1; u1; y1]

# objective

function objective(x, u)
    Q = Diagonal(ones(nx))
    q = ones(nx)

    R = Diagonal(0.1 * ones(nu))
    r = zeros(nu)

    return log(transpose(x) * Q * x) + exp(dot(q, x)) + sin(transpose(u) * R * u)
end

obj = objective(x, u)
obj_x = Symbolics.gradient(obj, x)
obj_u = Symbolics.gradient(obj, u)
obj_xu = Symbolics.gradient(obj, [x; u])

obj_xx = Symbolics.hessian(obj, x)
obj_uu = Symbolics.hessian(obj, u)
obj_xuxu = Symbolics.hessian(obj, [x; u])

obj_xx_sp = Symbolics.sparsehessian(obj, x)
obj_uu_sp = Symbolics.sparsehessian(obj, u)
obj_xuxu_sp = Symbolics.sparsehessian(obj, [x; u])
typeof(obj_func[1])
obj_func = Symbolics.build_function([obj], x, u)
obj_x_func = Symbolics.build_function(obj_x, x, u)
obj_u_func = Symbolics.build_function(obj_u, x, u)
obj_xu_func = Symbolics.build_function(obj_xu, x, u)

obj_xx_func = Symbolics.build_function(obj_xx_sp.nzval, x, u)
obj_uu_func = Symbolics.build_function(obj_uu_sp.nzval, x, u)
obj_xuxu_func = Symbolics.build_function(obj_xuxu_sp.nzval, x, u)

f = eval(obj_func[1])
@benchmark f($x1, $u1)
@benchmark objective($x1, $u1)

fx = eval(obj_x_func[1])
@benchmark fx($x1, $u1)
@benchmark ForwardDiff.gradient(a -> objective(a, $u1), $x1)

fu = eval(obj_u_func[1])
@benchmark fu($x1, $u1)
@benchmark ForwardDiff.gradient(a -> objective($x1, a), $u1)

fxu = eval(obj_xu_func[1])
@benchmark fxu($x1, $u1)

fxx = eval(obj_xx_func[1])
@benchmark fxx($x1, $u1)

fuu = eval(obj_uu_func[1])
@benchmark fuu($x1, $u1)

fxuxu = eval(obj_xuxu_func[1])
@benchmark fxuxu($x1, $u1)

obj_1 = zeros(1)
obj_x1 = zeros(nx)
obj_u1 = zeros(nu)
obj_xu1 = zeros(nx + nu)

obj_xx1 = zeros(length(obj_xx_sp.nzval))
obj_uu1 = zeros(length(obj_uu_sp.nzval))
obj_xuxu1 = zeros(length(obj_xuxu_sp.nzval))

f! = eval(obj_func[2])
@benchmark f!($obj_1, $x1, $u1)

fx! = eval(obj_x_func[2])
@benchmark fx!($obj_x1, $x1, $u1)
@benchmark ForwardDiff.gradient!($obj_x1, a -> objective(a, $u1), $x1)

fu! = eval(obj_u_func[2])
@benchmark fu!($obj_u1, $x1, $u1)
@benchmark ForwardDiff.gradient!($obj_u1, a -> objective($x1, a), $u1)

fxu! = eval(obj_xu_func[2])
@benchmark fxu!($obj_xu1, $x1, $u1)

fxx! = eval(obj_xx_func[2])
@benchmark fxx!($obj_xx1, $x1, $u1)
# @benchmark ForwardDiff.hessian!($obj_xx1, a -> objective(a, $u1), $x1)

fuu! = eval(obj_uu_func[2])
@benchmark fuu!($obj_uu1, $x1, $u1)
# @benchmark ForwardDiff.hessian!($obj_uu1, a -> objective($x1, a), $u1)

fxuxu! = eval(obj_xuxu_func[2])
@benchmark fxuxu!($obj_xuxu1, $x1, $u1)

# discrete dynamics
function quaternion_rotation_matrix(q)
	r, i, j, k  = q

	r11 = 1.0 - 2.0 * (j^2.0 + k^2.0)
	r12 = 2.0 * (i * j - k * r)
	r13 = 2.0 * (i * k + j * r)

	r21 = 2.0 * (i * j + k * r)
	r22 = 1.0 - 2.0 * (i^2.0 + k^2.0)
	r23 = 2.0 * (j * k - i * r)

	r31 = 2.0 * (i * k - j * r)
	r32 = 2.0 * (j * k + i * r)
	r33 = 1.0 - 2.0 * (i^2.0 + j^2.0)

	[r11 r12 r13;
     r21 r22 r23;
     r31 r32 r33]
end

function mrp_quaternion_map(mrp)
	n2 = mrp[1] * mrp[1] + mrp[2] * mrp[2] + mrp[3] * mrp[3]
    M = 2 / (1 + n2)
    return [(1 - n2) / (1 + n2), M * mrp[1], M * mrp[2], M * mrp[3]]
end

mrp_rotation_matrix(mrp) = quaternion_rotation_matrix(mrp_quaternion_map(mrp))

function dynamics(z, u, w)
#     mass = 1.0
#     lc = 1.0
#     gravity = 9.81
#     damping = 0.1

#     [z[2], (u[1] / ((mass * lc * lc)) - gravity * sin(z[1]) / lc - damping * z[2] / (mass * lc * lc))]
# end

    # states
    x = z[1:3]
    r = z[4:6]
    v = z[7:9]
    ω = z[10:12]

    # controls
    w1 = u[1]
    w2 = u[2]
    w3 = u[3]
    w4 = u[4]

    # forces
    kf = 0.1
    F1 = kf * w1
    F2 = kf * w2
    F3 = kf * w3
    F4 = kf * w4

    F = [0.0, 0.0, F1 + F2 + F3 + F4] # total rotor force in body frame

    # moments
    km = 0.01
    M1 = km * w1
    M2 = km * w2
    M3 = km * w3
    M4 = km * w4

    L = 0.25
    τ = [L * (F2 - F4), L * (F3 - F1), (M1 - M2 + M3 - M4)] # total rotor torque in body frame

    gravity = 9.81
    mass = 1.0
    inertia = Diagonal(ones(3))
    inertia_inv = Diagonal(ones(3))

    [v;
     0.25 * ((1.0 - r' * r) * ω - 2.0 * cross(ω, r) + 2.0*(ω' * r) * r);
     [0.0, 0.0, -9.81] + (1.0 / mass) * mrp_rotation_matrix(r) * F;
     inertia_inv * (τ - cross(ω, inertia * ω))]
end

h = 0.1
function discrete_dynamics(y, x, u, w)
    y - (x + h * dynamics(y, u, w))
end

dyn = discrete_dynamics(y, x, u, w)
dyn = simplify.(dyn)

# dense
dyn_y = Symbolics.jacobian(dyn, y)
dyn_x = Symbolics.jacobian(dyn, x)
dyn_u = Symbolics.jacobian(dyn, u)

dyn_xuy = Symbolics.jacobian(dyn, [x; u; y])

dyn_func = Symbolics.build_function(dyn, y, x, u, w)
dyn_y_func = Symbolics.build_function(dyn_y, y, x, u, w)
dyn_x_func = Symbolics.build_function(dyn_x, y, x, u, w)
dyn_u_func = Symbolics.build_function(dyn_u, y, x, u, w)
dyn_xuy_func = Symbolics.build_function(dyn_xuy, y, x, u, w)

# sparse
dyn_y_sp = Symbolics.sparsejacobian(dyn, y)
dyn_x_sp = Symbolics.sparsejacobian(dyn, x)
dyn_u_sp = Symbolics.sparsejacobian(dyn, u)
dyn_xuy_sp = Symbolics.sparsejacobian(dyn, [x; u; y])

dyn_y_func_sp = Symbolics.build_function(dyn_y_sp.nzval, y, x, u, w)
dyn_x_func_sp = Symbolics.build_function(dyn_x_sp.nzval, y, x, u, w)
dyn_u_func_sp = Symbolics.build_function(dyn_u_sp.nzval, y, x, u, w)
dyn_xuy_func_sp = Symbolics.build_function(dyn_xuy_sp.nzval, y, x, u, w)

d = eval(dyn_func[1])
@benchmark d($y1, $x1, $u1, $w1)
@benchmark discrete_dynamics($y1, $x1, $u1, $w1)

dy = eval(dyn_y_func[1])
@benchmark dy($y1, $x1, $u1, $w1)
@benchmark ForwardDiff.jacobian(a -> discrete_dynamics(a, $x1, $u1, $w1), $y1)
dy_sp = eval(dyn_y_func_sp[1])
@benchmark dy_sp($y1, $x1, $u1, $w1)

dx = eval(dyn_x_func[1])
@benchmark dx($y1, $x1, $u1, $w1)
@benchmark ForwardDiff.jacobian(a -> discrete_dynamics($y1, a, $u1, $w1), $x1)
dx_sp = eval(dyn_x_func_sp[1])
@benchmark dx_sp($y1, $x1, $u1, $w1)

du = eval(dyn_u_func[1])
@benchmark du($y1, $x1, $u1, $w1)
@benchmark ForwardDiff.jacobian(a -> discrete_dynamics($y1, $x1, a, $w1), $u1)
du_sp = eval(dyn_u_func_sp[1])
@benchmark du_sp($y1, $x1, $u1, $w1)

dxuy = eval(dyn_xuy_func[1])
@benchmark dxuy($y1, $x1, $u1, $w1)
dxuy_sp = eval(dyn_xuy_func_sp[1])
@benchmark dxuy_sp($y1, $x1, $u1, $w1)

dyn_1 = zeros(nx)
dyn_y1 = zeros(nx, nx)
dyn_x1 = zeros(nx, nx)
dyn_u1 = zeros(nx, nu)
dyn_xuy1 = zeros(nx, nx + nx + nu)

dyn_y1_sp = zeros(length(dyn_y_sp.nzval))
dyn_x1_sp = zeros(length(dyn_x_sp.nzval))
dyn_u1_sp = zeros(length(dyn_u_sp.nzval))
dyn_xuy1_sp = zeros(length(dyn_xuy_sp.nzval))

d! = eval(dyn_func[2])
@benchmark d!($dyn_1, $y1, $x1, $u1, $w1)

dy! = eval(dyn_y_func[2])
@benchmark dy!($dyn_y1, $y1, $x1, $u1, $w1)
@benchmark ForwardDiff.jacobian!($dyn_y1, a -> discrete_dynamics(a, $x1, $u1, $w1), $y1)
dy_sp! = eval(dyn_y_func_sp[2])
@benchmark dy_sp!($dyn_y1_sp, $y1, $x1, $u1, $w1)

dx! = eval(dyn_x_func[2])
@benchmark dx!($dyn_x1, $y1, $x1, $u1, $w1)
@benchmark ForwardDiff.jacobian!($dyn_x1, a -> discrete_dynamics($y1, a, $u1, $w1), $x1)
dx_sp! = eval(dyn_x_func_sp[2])
@benchmark dx_sp!($dyn_x1_sp, $y1, $x1, $u1, $w1)

du! = eval(dyn_u_func[2])
@benchmark du!($dyn_u1, $y1, $x1, $u1, $w1)
@benchmark ForwardDiff.jacobian!($dyn_u1, a -> discrete_dynamics($y1, $x1, a, $w1), $u1)
du_sp! = eval(dyn_u_func_sp[2])
@benchmark du_sp!($dyn_u1_sp, $y1, $x1, $u1, $w1)

dxuy! = eval(dyn_xuy_func[2])
@benchmark dxuy!($dyn_xuy1, $y1, $x1, $u1, $w1)
dxuy_sp! = eval(dyn_xuy_func_sp[2])
@benchmark dxuy_sp!($dyn_xuy1_sp, $y1, $x1, $u1, $w1)
dxuy_sp!(dyn_xuy1_sp, y1, x1, u1, w1)
dyn_xuy1_sp

@variables λ[1:nx]
dyn_λ = dot(λ, dyn)
dyn_λ_xuy = Symbolics.hessian(dyn_λ, [x; u; y])
dyn_λ_xuy_sp = Symbolics.sparsehessian(dyn_λ, [x; u; y])

dyn_λ_xuy_func = Symbolics.build_function(dyn_λ_xuy, y, x, u, w, λ)
dyn_λ_xuy_func_sp = Symbolics.build_function(dyn_λ_xuy_sp.nzval, y, x, u, w, λ)

λ1 = ones(nx)
dyn_λ_xuy1 = zeros(nx + nx + nu, nx + nx + nu)
dyn_λ_xuy1_sp = zeros(length(dyn_λ_xuy_sp.nzval))

dλ_xuy! = eval(dyn_λ_xuy_func[2])
@benchmark dλ_xuy!($dyn_λ_xuy1, $y1, $x1, $u1, $w1, $λ1)

dλ_xuy_sp! = eval(dyn_λ_xuy_func_sp[2])
@benchmark dλ_xuy_sp!($dyn_λ_xuy1_sp, $y1, $x1, $u1, $w1, $λ1)

# state constraint
function state_constraint(x)
    x - ones(nx)
end

# control constraint
function control_constraint(u)
    [u - ones(nu); ones(nu) - u]
end

# state control constraint
function state_control_constraint(x, u)

end

# indices
T = 10
nx_dim = [nx for t = 1:T]
nu_dim = [nu for t = 1:T-1]
nw_dim = [nw for t = 1:T-1]
nxu_dim = [[nx + nu for t = 1:T-1]..., nx]

idx_x = [collect((t > 1 ? sum(nx_dim[1:t-1]) + sum(nu_dim[1:t-1]) : 0) .+ (1:nx_dim[t])) for t = 1:T]
idx_u = [collect((t > 1 ? sum(nx_dim[1:t-1]) + sum(nu_dim[1:t-1]) : 0) + nx_dim[t] .+ (1:nu_dim[t])) for t = 1:T-1]
idx_xu = [collect((t > 1 ? sum(nx_dim[1:t-1]) + sum(nu_dim[1:t-1]) : 0) .+ (1:nx_dim[t] + nu_dim[t])) for t = 1:T-1]
idx_xuy = [collect((t > 1 ? sum(nx_dim[1:t-1]) + sum(nu_dim[1:t-1]) : 0) .+ (1:nx_dim[t] + nu_dim[t] + nx_dim[t+1])) for t = 1:T-1]

typeof(idx_x)
num_var = sum(nx_dim) + sum(nu_dim)
z1 = randn(num_var)

# objective
obj_grad = zeros(num_var)

xv = [z1[idx_x[t]] for t = 1:T]
uv = [[z1[idx_u[t]] for t = 1:T-1]..., zeros(0)]
wv = [zeros(nw) for t = 1:T-1]

objv = [0.0]
gradv = [[obj_grad[idx_xu[t]] for t = 1:T-1]..., obj_grad[idx_x[T]]]

function eval_obj!(J, x, u, T)
    for t = 1:T
        f!(J, x[t], u[t])
    end
    return J
end

eval_obj!(objv, xv, uv, T)
@benchmark eval_obj!($objv, $xv, $uv, $T)

function eval_obj_grad!(grad, x, u, T)
    for t = 1:T
        fxu!(grad[t], x[t], u[t])
    end
end

eval_obj_grad!(gradv, xv, uv, T)
@benchmark eval_obj_grad!($gradv, $xv, $uv, $T)

# discrete dynamics
nd_dim = nx_dim[2:T]
num_con = sum(nx_dim[2:T])
idx_dyn = [(t > 1 ? sum(nx_dim[1 .+ (1:t-1)]) : 0) .+ (1:nx_dim[t+1]) for t = 1:T-1]
con1 = zeros(num_con)
conv = [con1[idx_dyn[t]] for t = 1:T-1]

function eval_con!(con, x, u, w, T)
    for t = 1:T-1
        d!(con[t], x[t+1], x[t], u[t], w[t])
    end
end

d!(conv[t], xv[t+1], xv[t], uv[t], wv[t])
eval_con!(conv, xv, uv, wv, T)
@benchmark eval_con!($conv, $xv, $uv, $wv, $T)

function eval_con!(con, x, u, w, T)
    for t = 1:T-1
        d!(con[t], x[t+1], x[t], u[t], w[t])
    end
end

num_jac = length(dyn_xuy_sp.nzval) * (T - 1)
idx_jac = [(t - 1) * length(dyn_xuy_sp.nzval) .+ (1:length(dyn_xuy_sp.nzval)) for t = 1:T-1]
jac1 = zeros(num_jac)
jacv = [jac1[idx_jac[t]] for t = 1:T-1]

function eval_jac!(jac, x, u, w, T)
    for t = 1:T-1
        dxuy_sp!(jac[t], x[t+1], x[t], u[t], w[t])
    end
end

dxuy_sp!(jacv[t], xv[t+1], xv[t], uv[t], wv[t])
eval_jac!(jacv, xv, uv, wv, T)
@benchmark eval_jac!($jacv, $xv, $uv, $wv, $T)

I_jac, J_jac, V_jac = findnz(dyn_xuy_sp)

function sparsity_jac(jac, nd_dim, nxu_dim, T;
    row_shift = 0, col_shift = 0)

    row = Int[]
    col = Int[]

    I, J, V = findnz(jac)

    for t = 1:T-1
        push!(row, (I .+ row_shift)...)
        push!(col, (J .+ col_shift)...)

        row_shift += nd_dim[t]
        col_shift += nxu_dim[t]
    end

    return collect(zip(row, col))
end

sparsity_jac(dyn_xuy_sp, nd_dim, nxu_dim, T)


text:
""" Commutative diagrams in a category.

**NOTE**: This module is at an early stage of construction. At this time, only
commutative squares are implemented.
"""
module CommutativeDiagrams
export SquareDiagram, ob, hom, left, right, top, bottom

using AutoHashEquals

using ...GAT
using ...Theories: DoubleCategory
import ...Theories: ob, hom, dom, codom, compose, ⋅, ⋆, HomV, HomH,
  composeH, composeV, id, idH, idV, id2, id2H, id2V, left, right, top, bottom
using ..FinSets

# Commutative squares
#####################

"""    SquareDiagram(top, bottom, left, right)

Creates a square diagram in a category, which forms the 2-cells of the double
category Sq(C). The four 1-cells are given in top, bottom, left, right order, to
match the GAT of a double category.
"""
@auto_hash_equals struct SquareDiagram{Ob,Hom,Homs<:AbstractVector{Hom}}
  corners::Vector{Ob}
  sides::Homs
end

function SquareDiagram(homs::AbstractVector)
  length(homs) == 4 ||
    error("Square diagrams accept exactly 4 homs, in order top, bottom, left, right")
  obs = [dom(homs[1]), dom(homs[2]), dom(homs[4]), codom(homs[4])]

  # checking well-formedness
  # top-left share domains
  @assert dom(homs[1]) == dom(homs[3])
  # bottom-right share codomains
  @assert codom(homs[2]) == codom(homs[4])
  # left-bottom intersection
  @assert codom(homs[3]) == dom(homs[2])
  # top-right intersection
  @assert codom(homs[1]) == dom(homs[4])

  SquareDiagram(obs, homs)
end

SquareDiagram(top, bottom, left, right) =
  SquareDiagram([top, bottom, left, right])

ob(sq::SquareDiagram) = sq.corners
hom(sq::SquareDiagram) = sq.sides

top(sq::SquareDiagram) = sq.sides[1]
bottom(sq::SquareDiagram) = sq.sides[2]
left(sq::SquareDiagram) = sq.sides[3]
right(sq::SquareDiagram) = sq.sides[4]

""" Arrow category of FinSet as a double category of commutative squares.

For any category C, you can construct Sq(C) as a
[double category](https://ncatlab.org/nlab/show/double+category)
of commutative squares.

TODO: This construction has nothing specifically to do with finite sets and
functions and should be generalized to any category C.
"""
@instance DoubleCategory{FinSet, FinFunction, FinFunction, SquareDiagram} begin
  @import dom, codom, top, bottom, left, right, ⋅
  idH(A::FinSet) = FinFunction(identity, A, A)
  idV(A::FinSet) = FinFunction(identity, A, A)

  composeH(f::FinFunction, g::FinFunction) = compose(f,g)
  composeV(f::FinFunction, g::FinFunction) = compose(f,g)

  id2(A::FinSet) = SquareDiagram(idH(A), idH(A), idV(A), idV(A))
  # FIXME: how do you distinguish between vertical and horizontal if they are the same type?
  id2V(f::FinFunction) = SquareDiagram(f, f, idV(A), idV(A))
  id2H(f::FinFunction) = SquareDiagram(idH(A), idH(A), f, f)

  """    composeH(s₁, s₂)

  compose two squares horizontally as shown below:
      1   -f->   3  -g->   5
      |          |         |
      |          |         |
      v          v         v
      2  -f'->   4  -g'->  6
   """
  function composeH(s₁::SquareDiagram, s₂::SquareDiagram)
    @assert ob(s₁)[3] == ob(s₂)[1]
    @assert ob(s₁)[4] == ob(s₂)[2]
    @assert right(s₁) == left(s₂)

    f = top(s₁)
    f′= bottom(s₁)
    g = top(s₂)
    g′= bottom(s₂)
    return SquareDiagram(f⋅g, f′⋅g′, left(s₁), right(s₂))
  end

  """    composeV(s₁, s₂)

  compose two squares vertically as shown below:
      1   -->  3
      |        |
      |        |
      v        v
      2  -->   4
      |        |
      |        |
      v        v
      5  -->   6
  """
  function composeV(s₁::SquareDiagram, s₂::SquareDiagram)
    @assert ob(s₁)[2] == ob(s₂)[1]
    @assert ob(s₁)[4] == ob(s₂)[3]
    @assert bottom(s₁) == top(s₂)

    f = left(s₁)
    f′= right(s₁)
    g = left(s₂)
    g′= right(s₂)
    return SquareDiagram(top(s₁), bottom(s₂), f⋅g, f′⋅g′)
  end
end

end


text:
import Ipopt, NLopt, Percival
using Nonconvex, LinearAlgebra, Test, Zygote, FiniteDifferences
const FDM = FiniteDifferences

@testset "OrderedDict of reals" begin
    f(x) = sqrt(x[:b])
    g(x, a, b) = (a*x[:a] + b)^3 - x[:b]

    m = DictModel(f)
    addvar!(m, :a, 0.0, 10.0)
    addvar!(m, :b, 0.0, 10.0)
    add_ineq_constraint!(m, x -> g(x, 2, 0))
    add_ineq_constraint!(m, x -> g(x, -1, 1))

    @test getmin(m) == OrderedDict(:a => 0.0, :b => 0.0)
    @test getmax(m) == OrderedDict(:a => 10.0, :b => 10.0)
    @test getmin(m, :a) == 0.0
    @test getmax(m, :a) == 10.0
    @test getmin(m, :b) == 0.0
    @test getmax(m, :b) == 10.0

    p = OrderedDict(:a => 1.234, :b => 2.345)

    vec_model, _, un = Nonconvex.tovecmodel(m)
    pvec = Nonconvex.flatten(p)[1]
    val0, grad0 = Nonconvex.value_gradient(Nonconvex.getobjective(vec_model), pvec)
    @test val0 == f(un(pvec))
    @test grad0 == Zygote.gradient(x -> f(un(x)), pvec)[1]

    val1, grad1 = Nonconvex.value_gradient(Nonconvex.getineqconstraint(vec_model, 1), pvec)
    @test val1 == g(un(pvec), 2, 0)
    @test grad1 == Zygote.gradient(pvec -> g(un(pvec), 2, 0), pvec)[1]

    val2, grad2 = Nonconvex.value_gradient(Nonconvex.getineqconstraint(vec_model, 2), pvec)
    @test val2 == g(un(pvec), -1, 1)
    @test grad2 == Zygote.gradient(pvec -> g(un(pvec), -1, 1), pvec)[1]

    vals, jac = Nonconvex.value_jacobian(Nonconvex.getineqconstraints(vec_model), pvec)
    @test [val1, val2] == vals
    @test [grad1 grad2]' == jac

    vals, jac = Nonconvex.value_jacobian(Nonconvex.getobjectiveconstraints(vec_model), pvec)
    @test [val0, val1, val2] == vals
    @test [grad0 grad1 grad2]' == jac

    for alg in [MMA87(), MMA02(), IpoptAlg(), NLoptAlg(:LD_MMA), PercivalAlg()]
        r = optimize(m, MMA02(), p)
        @test norm([r.minimizer[:a] - 1/3, r.minimizer[:b] - 8/27]) < 1e-6
    end
end

struct S
    a
    b
end
@testset "Vector of structs" begin
    f(x) = sqrt(x[1].b)
    g(x, a, b) = (a*x[1].a + b)^3 - x[1].b

    m = Model(f)
    addvar!(m, S(0.0, 0.0), S(10.0, 10.0))
    add_ineq_constraint!(m, x -> g(x, 2, 0))
    add_ineq_constraint!(m, x -> g(x, -1, 1))

    p = [S(1.234, 2.345)]

    vec_model, _, un = Nonconvex.tovecmodel(m)
    pvec, un = Nonconvex.flatten(p)
    val0, grad0 = Nonconvex.value_gradient(Nonconvex.getobjective(vec_model), pvec)
    @test val0 == f(un(pvec))
    @test grad0 == Zygote.gradient(x -> f(un(x)), pvec)[1]

    val1, grad1 = Nonconvex.value_gradient(Nonconvex.getineqconstraint(vec_model, 1), pvec)
    @test val1 == g(un(pvec), 2, 0)
    @test grad1 == Zygote.gradient(pvec -> g(un(pvec), 2, 0), pvec)[1]

    val2, grad2 = Nonconvex.value_gradient(Nonconvex.getineqconstraint(vec_model, 2), pvec)
    @test val2 == g(un(pvec), -1, 1)
    @test grad2 == Zygote.gradient(pvec -> g(un(pvec), -1, 1), pvec)[1]

    vals, jac = Nonconvex.value_jacobian(Nonconvex.getineqconstraints(vec_model), pvec)
    @test [val1, val2] == vals
    @test [grad1 grad2]' == jac

    vals, jac = Nonconvex.value_jacobian(Nonconvex.getobjectiveconstraints(vec_model), pvec)
    @test [val0, val1, val2] == vals
    @test [grad0 grad1 grad2]' == jac

    for alg in [MMA87(), MMA02(), IpoptAlg(), NLoptAlg(:LD_MMA), PercivalAlg()]
        r = optimize(m, MMA02(), p)
        @test norm([r.minimizer[1].a - 1/3, r.minimizer[1].b - 8/27]) < 1e-6
    end
end


text:
module DocumentFormat
using CSTParser
using CSTParser.Tokenize
import CSTParser.Tokenize.Tokens
using CSTParser: headof, valof, EXPR
using FilePathsBase

const default_options = (4, true, true, true, true, true, true, true, true, false, true, "none")

struct FormatOptions
    indent::Integer
    indents::Bool
    ops::Bool
    tuples::Bool
    curly::Bool
    calls::Bool
    iterOps::Bool
    comments::Bool
    docs::Bool
    lineends::Bool
    keywords::Bool
    kwarg::String # Options arg-> "none", "single", "off"
end

# Provide sane default constructor for applications who don't know option specifics.
# Any nothing argument is given a sane default value.
FormatOptions() = FormatOptions(default_options...)
function FormatOptions(
        indent::Union{Nothing,Integer},
        indents::Union{Nothing,Bool},
        ops::Union{Nothing,Bool},
        tuples::Union{Nothing,Bool},
        curly::Union{Nothing,Bool},
        calls::Union{Nothing,Bool},
        iterOps::Union{Nothing,Bool},
        comments::Union{Nothing,Bool},
        docs::Union{Nothing,Bool},
        lineends::Union{Nothing,Bool},
        keywords::Union{Nothing,Bool},
        kwarg::Union{Nothing,String}
    )
    FormatOptions(
        something(indent, default_options[1]),
        something(indents, default_options[2]),
        something(ops, default_options[3]),
        something(tuples, default_options[4]),
        something(curly, default_options[5]),
        something(calls, default_options[6]),
        something(iterOps, default_options[7]),
        something(comments, default_options[8]),
        something(docs, default_options[9]),
        something(lineends, default_options[10]),
        something(keywords, default_options[11]),
        something(kwarg, default_options[12])
    )
end

struct Edit{T}
    loc::T
    text::String
end

mutable struct State{T}
    offset::Int
    edits::T
    opts::FormatOptions
    text
    lines::Vector{Tuple{Int,Int}}
end

function format(original_text::AbstractString, formatopts::FormatOptions=FormatOptions())
    text = deepcopy(original_text)
    original_ast = CSTParser.remlineinfo!(Meta.parse(string("begin\n", text, "\nend"), raise=false))
    if original_ast.head == :error
        @warn ("There was an error in the original ast, original text returned.")
        return text
    end
    state = State(0, Edit[], formatopts, text, get_lines(text))
    x = CSTParser.parse(text, true)
    if formatopts.ops
        pass(x, state, operator_pass)
    end
    if formatopts.tuples
        state.offset = 0
        pass(x, state, tuple_pass)
    end
    if formatopts.curly
        state.offset = 0
        pass(x, state, curly_pass)
    end
    if formatopts.calls
        state.offset = 0
        pass(x, state, call_pass)
    end
    if formatopts.iterOps
        state.offset = 0
        pass(x, state, forloop_pass)
    end
    if formatopts.comments
        comments_pass(text, state)
    end
    if formatopts.docs
        state.offset = 0
        pass(x, state, doc_pass)
    end
    if formatopts.keywords
        state.offset = 0
        pass(x, state, kw_pass)
    end

    if formatopts.lineends
        lineends_pass(text, x, state)
    end
    sort!(state.edits, lt=(a, b) -> first(a.loc) < first(b.loc), rev=true)

    for i = 1:length(state.edits)
        text = apply(text, state.edits[i])
    end
    if formatopts.indents
        text = indents(text, state.opts)
    end
    new_ast = CSTParser.remlineinfo!(Meta.parse(string("begin\n", text, "\nend"), raise=false))
    if new_ast.head == :error
        @warn ("There was an error in the formatted text, original returned.")
        return original_text
    elseif original_ast != new_ast
        @warn ("Mismatch between AST of original and formatted text")
        return original_text
    end

    return text
end

function format(path::AbstractPath, formatopts::FormatOptions=FormatOptions())
    if isfile(path)
        extension(path) != "jl" && error("Only .jl files can be formatted.")

        content = read(path, String)
        formatted_content = format(content, formatopts)
        write(path, formatted_content)
    elseif exists(path)
        for p in walkpath(path)
            if extension(p) == "jl"
                format(p, formatopts)
            end
        end
    else
        error("Invalid path.")
    end

    return nothing
end

function isformatted(text::AbstractString, formatopts::FormatOptions=FormatOptions())
    original_text = text
    new_text = format(text, formatopts)

    return original_text == new_text
end

function isformatted(path::AbstractPath, formatopts::FormatOptions=FormatOptions())
    if isfile(path)
        extension(path) != "jl" && error("Only .jl files can be formatted.")

        content = read(path, String)
        formatted_content = format(content, formatopts)

        return content == formatted_content
    elseif exists(path)
        for p in walkpath(path)
            if extension(p) == "jl"
                if !isformatted(p, formatopts)
                    return false
                end
            end
        end

        return true
    else
        error("Invalid path.")
    end

    return nothing
end

function pass(x, state, f=(x, state) -> nothing)
    f(x, state)
    if length(x) > 0
        for a in x
            pass(a, state, f)
        end
    else
        state.offset += x.fullspan
    end
    state
end

function ensure_single_space_after(x, state, offset)
    if x.fullspan == x.span
        push!(state.edits, Edit(offset + x.fullspan, " "))
    end
end

function ensure_no_space_after(x, state, offset)
    if x.fullspan != x.span
        push!(state.edits, Edit(offset .+ (x.span + 1:x.fullspan), ""))
    end
end

function ensure_exactly_single_space_after(x, state, offset)
    if x.fullspan !== x.span + 1
        push!(state.edits, Edit(offset .+ (x.span + 1:x.fullspan), " "))
    end
end

function apply(text, edit::Edit{Int})
    v = Vector{UInt8}(deepcopy(text))
    String(vcat(v[1:edit.loc], Vector{UInt8}(edit.text), v[edit.loc + 1:end]))
end
function apply(text, edit::Edit{UnitRange{Int}})
    v = Vector{UInt8}(deepcopy(text))
    String(vcat(v[1:first(edit.loc) - 1], Vector{UInt8}(edit.text), v[last(edit.loc) + 1:end]))
end
include("passes.jl")
include("indents.jl")
include("utils.jl")
end


text:
using InterpHalo, Plots

points = range(0, stop = 1, length = 201)
orbits = range(0, stop = 1, length = 100)
P       = zeros(201,6,100)

for (valE,indE) in enumerate(orbits)
    for (val,ind) in enumerate(points)
        P[val,:,valE] = intH(ind,indE,2)
    end
end
plotlyjs()
plot(P[:,1,1],P[:,2,1],P[:,3,1], leg = false)
for valE in 2:length(orbits)
    plot!(P[:,1,valE],P[:,2,valE],P[:,3,valE])
end

gui()


text:
Base.eltype(::Type{<:AbstractCache{T}}) where {T} = T
Base.eltype(::Type{<:AbstractRecycler{T}}) where {T} = T

function Recyclers.get!(factory, cache::AbstractCache{T}) where {T}
    y = Recyclers.maybeget!(cache)
    if y === nothing
        return factory()::T
    else
        return something(y)::T
    end
end

Recyclers.get!(recycler::AbstractRecycler{T}) where {T} =
    Recyclers.get!(factoryof(recycler), cacheof(recycler))::T

Recyclers.recycle!(recycler::AbstractRecycler{T}, object::T) where {T} =
    Recyclers.recycle!(cacheof(recycler), object)

function Recyclers.recycling!(body, factory, cache::AbstractCache)
    object = Recyclers.get!(factory, cache)
    y = body(object)
    Recyclers.recycle!(cache, object)
    return y
end

function Recyclers.recycling!(body, recycler::AbstractRecycler)
    object = Recyclers.get!(recycler)
    y = body(object)
    Recyclers.recycle!(recycler, object)
    return y
end

# With element type
# (There's always a specialization)
#=
function (::Type{Cache})(objects; kwargs...) where {T,Cache<:AbstractCache{T}}
    cache = Cache(; kwargs...)
    for x in objects
        Recyclers.recycle!(cache, x) || break
    end
    return cache
end
=#

# No element type
function (::Type{Cache})(objects; kwargs...) where {Cache<:AbstractCache}
    if !(Base.IteratorEltype(objects) isa Base.HasEltype)
        objects = identity.(objects)
        isempty(objects) && error("`objects` collection has no eltype and is empty")
    end
    return Cache{eltype(objects)}(objects; kwargs...)
end

Recyclers.factoryof(recycler::AbstractRecycler) = recycler.factory
Recyclers.cacheof(recycler::AbstractRecycler) = recycler.cache
cachetypeof(::Type{Union{}}) = error("unreachable")

maybeeltype(::Type{<:AbstractRecycler}) = nothing
maybeeltype(::Type{<:AbstractRecycler{T}}) where {T} = T

function (::Type{Recycler})(factory; kwargs...) where {Recycler<:AbstractRecycler}
    Cache = cachetypeof(Recycler)
    T = maybeeltype(Recycler)
    if T === nothing
        object = factory()
        return Recycler(factory, Cache((object,); kwargs...))::Recycler
    else
        return Recycler(factory, Cache{T}(; kwargs...))::Recycler
    end
end

function Recyclers.unsafe_init!(recycler::AbstractRecycler)
    Recyclers.unsafe_init!(Recyclers.cacheof(recycler))
    return recycler
end

function Recyclers.unsafe_destroy!(recycler::AbstractRecycler)
    Recyclers.unsafe_destroy!(Recyclers.cacheof(recycler))
    return recycler
end

function Recyclers.unsafe_empty!(recycler::AbstractRecycler)
    Recyclers.unsafe_empty!(Recyclers.cacheof(recycler))
    return recycler
end


text:
using StreamingTimeSeries
using Base.Test
using TimeZones

# EMAFeature
f = EMAFeature(1.0, DateTime())
@test valueat(f, DateTime(2015,8,7)) == 0.0
update!(f, DateTime(2014,12,1,1,1,1), 1.0)
@test abs(valueat(f, DateTime(2015,8,7)) - 1.0) < 0.00001
update!(f, DateTime(2014,12,1,1,1,1), 1.0)
@test abs(valueat(f, DateTime(2015,8,7)) - 1.0) < 0.00001

f = EMAFeature(1.0, ZonedDateTime(DateTime(), TimeZone("UTC")))
@test valueat(f, ZonedDateTime(DateTime(2015,8,7), TimeZone("UTC"))) == 0.0
update!(f, ZonedDateTime(DateTime(2014,12,1,1,1,1), TimeZone("UTC")), 1.0)
@test abs(valueat(f, ZonedDateTime(DateTime(2015,8,7), TimeZone("UTC"))) - 1.0) < 0.00001
update!(f, ZonedDateTime(DateTime(2014,12,1,1,1,1), TimeZone("UTC")), 1.0)
@test abs(valueat(f, ZonedDateTime(DateTime(2015,8,7), TimeZone("UTC"))) - 1.0) < 0.00001

# EMVFeature
f = EMVFeature(2.0, DateTime())
@test valueat(f, DateTime(2015,8,7)) == 0.0
update!(f, DateTime(2015,8,15,9,0,0), 1.0)
update!(f, DateTime(2015,8,15,10,0,0), 10.0)
update!(f, DateTime(2015,8,15,11,0,0), 1.0)
update!(f, DateTime(2015,8,15,12,0,0), 10.0)
@test valueat(f, DateTime(2015,8,15,12,0,0)) > 0.1
update!(f, DateTime(2015,8,15,12,0,0), 2.0)
update!(f, DateTime(2015,8,15,13,0,0), 2.0)
update!(f, DateTime(2015,8,15,14,0,0), 2.0)
update!(f, DateTime(2015,8,15,15,0,0), 2.0)
update!(f, DateTime(2015,8,15,16,0,0), 2.0)
update!(f, DateTime(2015,8,15,17,0,0), 2.0)
update!(f, DateTime(2015,8,15,18,0,0), 2.0)
@test valueat(f, DateTime(2015,8,15,18,0,0)) < 0.1

# TimeSinceFeature
f = TimeSinceFeature(DateTime(2015,8,7,1,1,1))
@test valueat(f, DateTime(2015,8,7,1,1,1)) == 0.0
update!(f, DateTime(2015,8,7,2,1,1))
@test valueat(f, DateTime(2015,8,7,2,1,1)) == 0.0
@test valueat(f, DateTime(2015,8,7,2,3,1)) == 2.0
@test valueat(f, DateTime(2014,8,7,2,3,1)) == 0.0
f = TimeSinceFeature(DateTime())
update!(f, DateTime(2015,8,7,2,1,1))
@test valueat(f, DateTime(2015,8,7,2,1,1)) == 0.0
@test valueat(f, DateTime(2015,8,7,2,3,1)) == 2.0

# LastValueFeature
f = LastValueFeature(DateTime(2015,8,7), 1.4)
@test valueat(f, DateTime(2015,8,6)) == 0
@test valueat(f, DateTime(2015,8,7)) == 1.4
@test valueat(f, DateTime(2015,8,8)) == 1.4
f = LastValueFeature(DateTime())
@test valueat(f, DateTime(2015,8,6)) == 0

# DecayFeature
f = DecayFeature(1.0, DateTime())
@test valueat(f, DateTime(2015,8,7)) == 0.0
update!(f, DateTime(2014,12,1,1,1,1), 1.0)
@test valueat(f, DateTime(2015,8,7)) < 0.00001
update!(f, DateTime(2014,12,1,1,1,1), 1.0)
@test abs(valueat(f, DateTime(2014,12,1,1,2,1)) - 1) < 0.00001

# LagSumFeature
f = LagSumFeature(Dates.Minute(1), Dates.Minute(0), DateTime())
@test valueat(f, DateTime(2015,8,7)) == 0.0
update!(f, DateTime(2015,9,1,1,1,1), 1.0)
@test valueat(f, DateTime(2015,9,1,1,1,1)) == 1.0
update!(f, DateTime(2015,9,1,1,2,1), 1.0)
@test valueat(f, DateTime(2015,9,1,1,2,1)) == 1.0
@test valueat(f, DateTime(2015,9,1,1,3,1)) == 0.0

f = LagSumFeature(Dates.Minute(6), Dates.Minute(4), DateTime())
update!(f, DateTime(2015,9,1,1,1,1), 1.0)
update!(f, DateTime(2015,9,1,1,2,1), 1.0)
@test valueat(f, DateTime(2015,9,1,1,3,1)) == 0.0
@test valueat(f, DateTime(2015,9,1,1,5,1)) == 1.0
@test valueat(f, DateTime(2015,9,1,1,6,1)) == 2.0
@test valueat(f, DateTime(2015,9,1,1,7,1)) == 1.0
@test valueat(f, DateTime(2015,9,1,1,8,1)) == 0.0

# LagLastValueFeature
f = LagLastValueFeature(Dates.Second(1), DateTime())
@test valueat(f, DateTime(2015,8,7)) == 0.0
update!(f, DateTime(2015,9,1,1,1,1), 1.0)
@test valueat(f, DateTime(2015,9,1,1,1,1)) == 0.0
update!(f, DateTime(2015,9,1,1,1,2), 2.0)
@test valueat(f, DateTime(2015,9,1,1,1,2)) == 1.0
@test valueat(f, DateTime(2015,9,1,1,1,3)) == 2.0
@test valueat(f, DateTime(2015,9,1,1,1,4)) == 2.0


text:
using Test
using NiLang, NiLang.AD
using LinearAlgebra
using Random

@testset "inv" begin
    Random.seed!(2)

    id = [1 0 0; 0 1 0; 0 0 1.0]
    @test check_inv(i_inv!, (randn(3, 3), randn(3, 3)))
    @test check_inv(PlusEq(det), (0.3, randn(3, 3)))
    @test check_inv(PlusEq(logdet), (0.3, rand(3, 3) .+ id))
    @test check_grad(PlusEq(det), (0.3, randn(3, 3)), iloss=1)
    @test check_grad(PlusEq(logdet), (0.3, rand(3, 3) .+ id), iloss=1)

    @i function loss(out!, y, A)
        i_inv!(y, A)
        out! += y[1,1]
    end
    @test check_grad(loss, (0.0, randn(3, 3), randn(3, 3)); iloss=1)
end

@testset "affine" begin
    Random.seed!(2)
    A = randn(5, 5)
    b = randn(5)
    x = randn(5)
    y! = zeros(5)
    @test i_affine!(y!, A, b, x)[1] ≈ A*x + b
end