text:
using MendelIHT
using SnpArrays
using DataFrames
using Distributions
using BenchmarkTools
using Random
using LinearAlgebra
using StatsFuns: logistic

function time_logistic_response(
    n :: Int64,        # number of cases
    p :: Int64,        # number of predictors (SNPs)
    k :: Int64,        # number of true predictors per group
    debias :: Bool
)
	Random.seed!(1111)

	#construct snpmatrix, covariate files, and true model b
	x, maf = simulate_random_snparray(n, p)
	xbm = SnpBitMatrix{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true);
	z = ones(n, 1) # non-genetic covariates, just the intercept
	true_b = zeros(p)
	true_b[1:k] = randn(k)
	shuffle!(true_b)
	correct_position = findall(x -> x != 0, true_b)

	#simulate bernoulli data
	y_temp = xbm * true_b
	prob = logistic.(y_temp) #inverse logit link
	y = [rand(Bernoulli(x)) for x in prob]
	y = Float64.(y)

    #free memory from xbm.... does this work like I think it works?
	xbm = nothing
	GC.gc()

	#time the result and return
    result = @timed L0_logistic_reg(x, z, y, 1, k, glm = "logistic", debias=debias, show_info=false, convg=true, init=false)
    iter = result[1].iter
    runtime = result[2]      # seconds
    memory = result[3] / 1e6 # MB

    return iter, runtime, memory
end

function main(n :: Int64, i :: Int64)
	p = 1000
	k = 10
	debias = false
	response = "logistic"

	#run result: first run on super small data to force compile, then run result
	iter, runtime, memory = time_logistic_response(100, 100, k, debias)
	iter, runtime, memory = time_logistic_response(n, p, k, debias)

	#write result to corresponding file
	if debias
		dir = response * "_results/"
	else
		dir = response * "_results_nodebias/"
	end
	file = string(n) * "_by_" * string(p) * "_run$i"
	open(dir * file, "w") do f
		write(f, "time (seconds), memory (MB), iteration\n")
		write(f, string(runtime) * ", " * string(memory) * ", " * string(iter))
	end
end

# arguments passed from command line
n = parse(Int64, ARGS[1])
i = parse(Int64, ARGS[2]) #i'th run: time result 5 times so we can average (didn't use benchmarktools cuz takes too long)

main(n, i)


text:
function create_manifest(stream,
                         prefix :: AbstractString,
                         amount :: Integer,
                         ratio  :: Rational)
    dict = Dict("prefix"      => prefix,
                "amount"      => amount,
                "numerator"   => numerator(ratio),
                "denominator" => denominator(ratio))
    JSON.print(stream, dict, 4)
end

function create_archive(outname :: AbstractString,
                        slices  :: AbstractArray{Bool, 3},
                        ratio   :: Rational)
    prefix = slice_prefix(outname)
    amount = size(slices, 3)

    ZipFile.Writer(outname) do root
        manifest = ZipFile.addfile(root, datafile("manifest.json"), method = ZipFile.Deflate)
        create_manifest(manifest, prefix, amount, ratio)

        for k in 1:amount
            slice = slices[:,:,k]
            name  = datafile("$(prefix)-$(k).png")
            pic   = ZipFile.addfile(root, name, method = ZipFile.Deflate)
            # Here use another IO stream because ZipFile does not work
            # well with FileIO. Julia cannot into abstractions.
            tmpio = IOBuffer()
            save(Stream{format"PNG"}(tmpio), Gray.(slice))
            write(pic, take!(tmpio))
        end
    end

    return
end

function var"archive-w/-nth"(data    :: AbstractArray{Bool, 3},
                             outname :: AbstractString,
                             n       :: Int)
    @assert iszero(rem(size(data, 3), n))
    slices = data[:,:,1:n:end]

    return create_archive(outname, slices, n//1)
end

function var"archive-w/o-nth"(data    :: AbstractArray{Bool, 3},
                              outname :: AbstractString,
                              n       :: Int)
    @assert iszero(rem(size(data, 3), n))
    orig_ratio = n // (n - 1)
    slices = reduce(stich, data[:,:,k] for k in 1:size(data,3) if !iszero(mod(k, n)))

    return create_archive(outname, slices, orig_ratio)
end

"""
    archive(data, outname, n, mode)

Create an archive of slices from 3D bit array with two-phase porous
medium. When `mode` is `:take` (resp. `drop`) the archive is created
by taking (resp. dropping) each `n`-th sample from `data`. It's when
stored on disk with the name `outname`. This array has sufficient
information to recover the original data (with some loss).

See also: [`reconstruct`](@ref).
"""
function archive end

archive(data, outname, n, :: Val{:take}) =
    var"archive-w/-nth"(data, outname, n)

archive(data, outname, n, :: Val{:drop}) =
    var"archive-w/o-nth"(data, outname, n)

archive(data, outname, n, mode :: Symbol) =
    archive(data, outname, n, Val(mode))

ext: jl


text:
export
    Primals

function PartedArrays.create_partition(n::Int,m::Int,N::Int,uN=N-1)
    Nx = N*n
    Nu = uN*m
    Nz = Nx+Nu
    ind_x = zeros(Int,n,N)
    ind_u = zeros(Int,m,uN)
    ix = 1:n
    iu = n .+ (1:m)
    for k = 1:uN
        ind_x[:,k] = ix .+ (k-1)*(n+m)
        ind_u[:,k] = iu .+ (k-1)*(n+m)
    end
    if uN == N-1
        ind_x[:,N] = ix .+ (N-1)*(n+m)
    end
    return (X=ind_x, U=ind_u)
end

# TODO: Inherit from AbstractArray
struct Primals{T<:Real}
    Z::Vector{T}
    X::Vector{SubArray{T,1,Vector{T},Tuple{Vector{Int}},false}}
    U::Vector{SubArray{T,1,Vector{T},Tuple{Vector{Int}},false}}
    equal::Bool
end

function Primals(Z::Vector{T},n::Int,m::Int) where T
    if length(Z) % (n+m) == 0
        N = length(Z) ÷ (n+m)
        uN = N
    elseif length(Z) % (n+m) == n
        N = length(Z) ÷ (n+m) + 1
        uN = N-1
    end
    part_z = create_partition(n,m,N,uN)
    Primals(Z,part_z)
end

function Primals(Z,part_z::NamedTuple)
    N, uN = size(part_z.X,2), size(part_z.U,2)
    X = [view(Z,part_z.X[:,k]) for k = 1:N]
    U = [view(Z,part_z.U[:,k]) for k = 1:uN]
    Primals(Z,X,U, N==uN)
end

"""Create a Primals from vectors of SubArrays.
This is the fastest method to convert a vector Z to subarrays X,U. This will overwrite X and U!
"""
function Primals(Z::Vector{T},X::Vector{S},U::Vector{S}) where {T,S<:SubArray}
    N = length(X)
    uN = length(U)
    for k = 1:uN
        X[k] = view(Z,X[k].indices[1])
        U[k] = view(Z,U[k].indices[1])
    end
    if uN == N-1
        X[N] = view(Z,X[N].indices)
    end
    Primals(Z,X,U, N==uN)
end

function Primals(Z::Vector{T},Z0::Primals{T}) where T
    X = deepcopy(Z0.X)
    U = deepcopy(Z0.U)
    uN = length(U)
    N = length(X)
    for k = 1:uN
        X[k] = view(Z,X[k].indices[1])
        U[k] = view(Z,U[k].indices[1])
    end
    if uN == N-1
        X[N] = view(Z,X[N].indices)
    end
    Primals(Z,X,U, Z0.equal)
end

""" $(TYPEDSIGNATURES)
Combine state and control trajectories.
This is a little slow and less memory efficient than converting from a the combined vector to individual trajectories.
"""
function Primals(X::VectorTrajectory{T}, U::VectorTrajectory{T}) where T
    N,uN = length(X), length(U)
    n,m = length(X[1]), length(U[1])
    NN = N*n + uN*m
    Z = zeros(T,NN)
    part_z = create_partition(n,m,N,uN)
    Z[part_z.X] = to_array(X)
    Z[part_z.U] = to_array(U)
    Primals(Z,part_z)
end

function Primals(prob::Problem{T}, equal::Bool=false) where T
    U = copy(prob.U)
    if equal
        U = push!(U, U[end])
    end
    Primals(prob.X, U)
end

Base.size(Z::Primals) = length(Z.X[1]), length(Z.U[1]), length(Z.X)
Base.length(Z::Primals) = length(Z.Z)
Base.copy(Z::Primals) = Primals(copy(Z.Z),Z)

function pack(prob::Problem{T}) where T
    n,m,N = size(prob)
    part_z = create_partition(n,m,N,N)
    NN = N*(n+m)
    Z = PartedVector(zeros(T,NN),part_z)
    copyto!(Z.X, prob.X)
    copyto!(Z.U, prob.U)
    return Z
end

function unpack(Z::Vector{<:Real}, part_z::NamedTuple)
    N, uN = size(part_z.X,2), size(part_z.U,2)
    X = [view(Z,part_z.X[:,k]) for k = 1:N]
    U = [view(Z,part_z.U[:,k]) for k = 1:uN]
    return X, U
end

function pack(X,U, part_z)
    n,m = length(X[1]), length(U[1])
    N, uN = length(X), length(U)
    Z = zeros(eltype(X[1]), N*n + uN*m)
    for k = 1:N
        Z[part_z.X[:,k]] = X[k]
        Z[part_z.U[:,k]] = U[k]
    end
    return Z
end

function pack(X,U)
    n = length(X[1])
    m = length(U[1])
    N = length(X)
    uN = length(U)
    part_z = create_partition(n,m,N,uN)
    Z = pack(X,U,part_z)
end

struct PrimalDual{T} <: AbstractArray{T,1}
    V::Vector{T}
    Z::SubArray{T,1,Vector{T},Tuple{UnitRange{Int}},true}
    X::Vector{SubArray{T,1,Vector{T},Tuple{Vector{Int}},false}}
    U::Vector{SubArray{T,1,Vector{T},Tuple{Vector{Int}},false}}
    Y::SubArray{T,1,Vector{T},Tuple{UnitRange{Int}},true}
    ν::Vector{SubArray{T,1,Vector{T},Tuple{UnitRange{Int}},true}}
    λ::Vector{SubArray{T,1,Vector{T},Tuple{UnitRange{Int}},true}}
    active_set::Vector{Bool}
    a::Vector{SubArray{Bool,1,Vector{Bool},Tuple{UnitRange{Int}},true}}
end

function PrimalDual(prob::Problem{T}) where T
    n,m,N = size(prob)
    NN = N*n + (N-1)*m
    p_colloc = N*n
    p = num_constraints(prob)
    pcum = insert!(cumsum(p),1,0)
    P = sum(p) + p_colloc

    V = zeros(T, NN+P)
    Z = view(V,1:NN)
    Y = view(V,NN .+ (1:P))

    part_z = create_partition(n,m,N)
    X = [view(V,part_z.X[:,k]) for k = 1:N]
    U = [view(V,part_z.U[:,k]) for k = 1:N-1]

    # Constraints
    y_part = ones(Int,2,N-1)*n
    y_part[2,:] = p[1:end-1]
    y_part = vec(y_part)
    insert!(y_part,1,n)
    push!(y_part, p[N])
    y_inds = cumsum(y_part)
    insert!(y_inds,1,0)
    y_inds = [y_inds[k]+1:y_inds[k+1] for k = 1:length(y_inds)-1]
    ind_colloc = y_inds[insert!(collect(2:2:2N-1),1,1)]
    ind_custom = y_inds[push!(collect(3:2:2N),2N)]

    ν = [view(V, NN .+ inds) for inds in ind_colloc]
    λ = [view(V, NN .+ inds) for inds in ind_custom]
    active_set = ones(Bool,P)
    a = [view(active_set, p_colloc + pcum[k] .+ (1:p[k])) for k = 1:N]
    copyto!.(X, prob.X)
    copyto!.(U, prob.U)
    PrimalDual(V, Z, X, U, Y, ν, λ, active_set, a)
end

function PrimalDual(V::Vector, n::Int, m::Int, p::Vector{Int}, N::Int) where T
    NN = N*n + (N-1)*m
    p_colloc = N*n
    pcum = insert!(cumsum(p),1,0)
    P = sum(p) + p_colloc

    Z = view(V,1:NN)
    Y = view(V,NN .+ (1:P))

    part_z = create_partition(n,m,N)
    X = [view(V,part_z.X[:,k]) for k = 1:N]
    U = [view(V,part_z.U[:,k]) for k = 1:N-1]
    ν = [view(V, NN + (k-1)*n .+ (1:n)) for k = 1:N]
    λ = [view(V, NN + p_colloc + pcum[k] .+ (1:p[k])) for k = 1:N]
    active_set = ones(Bool,P)
    a = [view(active_set, p_colloc + pcum[k] .+ (1:p[k])) for k = 1:N]
    PrimalDual(V, Z, X, U, Y, ν, λ, active_set, a)
end
PrimalDual(V::Vector, prob::Problem) = PrimalDual(V, prob.model.n, prob.model.m, num_constraints(prob), prob.N)

function Base.copy(V::PrimalDual)
    V2 = copy(V.V)
    Z = view(V2, V.Z.indices[1])
    N = length(V.X)
    X = [view(Z, V.X[k].indices[1]) for k = 1:N]
    U = [view(Z, V.U[k].indices[1]) for k = 1:N-1]
    Y = view(V2, V.Y.indices[1])
    ν = [view(V2, V.ν[k].indices[1]) for k = 1:N]
    λ = [view(V2, V.λ[k].indices[1]) for k = 1:N]
    active_set = copy(V.active_set)
    a = [view(active_set, V.a[k].indices[1]) for k = 1:N]
    PrimalDual(V2, Z, X, U, Y, ν, λ, active_set, a)
end

import Base: size, length, getindex, setindex
size(V::PrimalDual) = size(V.V)
length(V::PrimalDual) = length(V.V)
getindex(V::PrimalDual, i::Int) = V.V[i]
setindex!(V::PrimalDual, v, i::Int) = setindex!(V.V, v, i)
IndexStyle(::PrimalDual) = IndexLinear()

+(V::PrimalDual, A::Vector) = begin V_ = copy(V); V_.V .+= A; V_ end

primals(V::PrimalDual) = V.Z
duals(V::PrimalDual) = V.Y

max_issues_repo_name: MehrdadZareian/GMARA

max_forks_repo_name: MehrdadZareian/GMARA


text:
export emsd_f, emsd_f_circuit, ρ_junct, ρ_junct_avg, ρ_junct_mean_std

using Statistics: mean, var, std

"""
```julia
emsd_f(f, params; N=32)
```
Computes the mean and standard deviation of an ensemble of `StickNetwork`s determined by `params` with ensemble size `N`.
"""
function emsd_f(f, params; N=32)
    e = ensemble(N, params...)
    res = f.(e)
    mean(res), std(res)
end

"""
```julia
emsd_f_circuit(f, params, M, args...; N=32)
```
Computes the mean and standard deviation of an ensemble of `N` nanowire networks
with model `M` and parameters 'params'. `args` may specify additonal parameters to
`ensemble_circuit` such as the `Val(:conducting)` property.
"""
function emsd_f_circuit(f, params, M, args...; N=32)
    e = ensemble_circuit(N, params..., M, args...)
    res = f.(e)
    mean(res), std(res)
end

"""
```julia
ρ_junct(net)
```
Computes junction density of `StickNetwork` `net`.
"""
function ρ_junct(net::StickNetwork{T,N}) where {T,N}
    n = length(net.lines)
    m = 0
    for i ∈ 2:(n)
        for j ∈ 1:(i-1)
            ps = intersection_params(net.lines[i],net.lines[j])
            if (0.0 ≤ ps[1] ≤ 1.0) && (0.0 ≤ ps[2] ≤ 1.0)
                p = line_segment_point(ps[1],net.lines[i])
                if sum( zeros(SVector{N,T}) .≤ p .≤ net.dims ) == N
                    m += 1
                end
            end
        end
    end
    m/prod(net.dims)
end

"""
```julia
ρ_junct_mean_std(params; N=32)
```
Computes the mean and standard deviation of the junction
density of an ensemble of `StickNetwork`s given by `params`.
"""
ρ_junct_mean_std(params; N=32) = emsd_f(ρ_junct, params; N=N)
ρ_junct_mean_std(params...; N=32) = ρ_junct_mean_std(params; N=N)

"""
Deprecated, use `ρ_junct_mean_std` instead.
"""
ρ_junct_avg(dims, ρₙ, l, ρ, D; N=32) = emsd_f(ρ_junct, [dims, ρₙ, l, ρ, D], N=N)[1]


text:
# # Double Gyre
#
# This example simulates a double gyre following:
# https://mitgcm.readthedocs.io/en/latest/examples/baroclinic_gyre/baroclinic_gyre.html

using LinearAlgebra
using Statistics
using Printf

using Oceananigans
using Oceananigans.Grids
using Oceananigans.Fields
using Oceananigans.Advection
using Oceananigans.AbstractOperations
using Oceananigans.Diagnostics
using Oceananigans.OutputWriters
using Oceananigans.Utils

using Oceananigans.Simulations: get_Δt

## Convective adjustment

function convective_adjustment!(model, Δt, K)
    grid = model.grid
    Nx, Ny, Nz = grid.Nx, grid.Ny, grid.Nz
    Δz = model.grid.Δz
    T = model.tracers.T

    ∂T∂z = ComputedField(@at (Center, Center, Center) ∂z(T))
    compute!(∂T∂z)

    κ = zeros(Nx, Ny, Nz)
    for k in 1:Nz, j in 1:Ny, i in 1:Nx
        κ[i, j, k] = ∂T∂z[i, j, k] < 0 ? K : 0
    end

    T_interior = interior(T)
    Tⁿ⁺¹ = zeros(Nx, Ny, Nz)

    for j in 1:Ny, i in 1:Nx
        ld = [-Δt/Δz^2 * κ[i, j, k]   for k in 2:Nz]
        ud = [-Δt/Δz^2 * κ[i, j, k+1] for k in 1:Nz-1]

        d = zeros(Nz)
        for k in 1:Nz-1
            d[k] = 1 + Δt/Δz^2 * (κ[i, j, k] + κ[i, j, k+1])
        end
        d[Nz] = 1 + Δt/Δz^2 * κ[i, j, Nz]

        𝓛 = Tridiagonal(ld, d, ud)

        Tⁿ⁺¹[i, j, :] .= 𝓛 \ T_interior[i, j, :]
    end

    set!(model, T=Tⁿ⁺¹)

    return nothing
end

## Grid setup

@info "Grid setup..."

km = kilometers
topo = (Bounded, Bounded, Bounded)
domain = (x=(-2000km, 2000km), y=(-3000km, 3000km), z=(-2km, 0))
grid = RegularCartesianGrid(topology=topo, size=(96, 96, 32); domain...)

## Boundary conditions

@info "Boundary conditions setup..."

@inline wind_stress(x, y, t, p) = - p.τ * cos(2π * y / p.L)
@inline u_bottom_stress(x, y, t, u, p) = - p.μ * p.H * u
@inline v_bottom_stress(x, y, t, v, p) = - p.μ * p.H * v

wind_stress_params = (τ=1e-4, L=grid.Ly)
wind_stress_bc = FluxBoundaryCondition(wind_stress, parameters=wind_stress_params)

bottom_stress_params = (μ=1/30day, H=grid.Lz)
u_bottom_stress_bc = FluxBoundaryCondition(u_bottom_stress, field_dependencies=:u, parameters=bottom_stress_params)
v_bottom_stress_bc = FluxBoundaryCondition(v_bottom_stress, field_dependencies=:v, parameters=bottom_stress_params)

no_slip = ValueBoundaryCondition(0)

u_bcs = UVelocityBoundaryConditions(grid,
       top = wind_stress_bc,
    bottom = u_bottom_stress_bc,
     north = no_slip,
     south = no_slip
)

v_bcs = VVelocityBoundaryConditions(grid,
      east = no_slip,
      west = no_slip,
    bottom = v_bottom_stress_bc
)

w_bcs = WVelocityBoundaryConditions(grid,
    north = no_slip,
    south = no_slip,
     east = no_slip,
     west = no_slip
)

@inline T_reference(y, p) = p.T_mid + p.ΔT / p.Ly * y
@inline temperature_flux(x, y, t, T, p) = @inbounds - p.μ * (T - T_reference(y, p))

T_min, T_max = 0, 30
temperature_flux_params = (T_min=T_min, T_max=T_max, T_mid=(T_min+T_max)/2, ΔT=T_max-T_min, μ=1/day, Ly=grid.Ly)
temperature_flux_bc = FluxBoundaryCondition(temperature_flux, field_dependencies=:T, parameters=temperature_flux_params)

T_bcs = TracerBoundaryConditions(grid,
    bottom = ValueBoundaryCondition(T_min),
       top = temperature_flux_bc
)

## Turbulent diffusivity closure

closure = AnisotropicDiffusivity(νh=500, νz=1e-2, κh=100, κz=1e-2)

## Model setup

@info "Model setup..."

model = IncompressibleModel(
           architecture = CPU(),
                   grid = grid,
            timestepper = :RungeKutta3,
              advection = WENO5(),
               coriolis = BetaPlane(latitude=45),
               buoyancy = SeawaterBuoyancy(constant_salinity=true),
                tracers = :T,
                closure = closure,
    boundary_conditions = (u=u_bcs, v=v_bcs, w=w_bcs, T=T_bcs)
)

## Initial condition

@info "Setting initial conditions..."

# a stable density gradient with random noise superposed.
T₀(x, y, z) = temperature_flux_params.T_min + temperature_flux_params.ΔT/2 * (1 + z / grid.Lz)
set!(model, T=T₀)

# set!(model, T=temperature_flux_params.T_mid)

## Simulation setup

@info "Setting up simulation..."

u_max = FieldMaximum(abs, model.velocities.u)
v_max = FieldMaximum(abs, model.velocities.v)
w_max = FieldMaximum(abs, model.velocities.w)

wall_clock = time_ns()

function print_progress(simulation)
    model = simulation.model

    K = 10
    convective_adjustment!(model, get_Δt(simulation), K)

    T_interior = interior(model.tracers.T)
    T_min, T_max = extrema(T_interior)
    T_mean = mean(T_interior)

    ## Print a progress message
    msg = @sprintf("i: %04d, t: %s, Δt: %s, u_max = (%.1e, %.1e, %.1e) m/s, T: (min=%.2f, mean=%.2f, max=%.2f), wall time: %s\n",
                   model.clock.iteration,
                   prettytime(model.clock.time),
                   prettytime(wizard.Δt),
                   u_max(), v_max(), w_max(), T_min, T_mean, T_max,
                   prettytime(1e-9 * (time_ns() - wall_clock)))

    @info msg

    return nothing
end

wizard = TimeStepWizard(cfl=0.5, diffusive_cfl=0.5, Δt=1hour, max_change=1.1, max_Δt=1hour)

simulation = Simulation(model, Δt=wizard, stop_time=2years, iteration_interval=1, progress=print_progress)

## Set up output writers

@info "Setting up output writers..."

simulation.output_writers[:fields] =
    NetCDFOutputWriter(model, merge(model.velocities, model.tracers),
                       schedule=TimeInterval(1day), filepath="double_gyre.nc", mode="c")

## Running the simulation

@info "Running simulation..."
run!(simulation)


text:
"""
    $(SIGNATURES)

Calculate ΔG'⁰ for a multicompartment reaction. Assume that `equilibrator` corresponds to
the thermodynamic conditions of the inner reaction. The `potential_difference` across
the membrane, as well as the `pH_outer`, `pMg_outer`, and `ionic_strength_outer` can be set.
The `tolerance` sets the allowed imbalance between the `rxn_inner` and `rxn_outer`.

# Example
```
equilibrator = Equilibrator()

cyto_rxn = bigg("adp + pi + 2 h = h2o + atp")
peri_rxn = bigg(" = 2 h")

ψ = 0.17u"V"
peri_ionic = 0.2u"M"
peri_pMg = 3.0
peri_pH = 6.5

set_ionic_strength(eq, 250u"mM")
set_pH(eq, 7.4)
set_pMg(eq, 3.0)

dg = multi_compartmental_standard_dg_prime(
    eq,
    cyto_rxn,
    peri_rxn;
    e_potential_difference = ψ,
    pH_outer = peri_pH,
    pMg_outer = peri_pMg,
    ionic_strength_outer = peri_ionic,
    tolerance=2.0,
)
```
"""
function multi_compartmental_standard_dg_prime(
    equilibrator::Equilibrator,
    rxn_inner::String,
    rxn_outer::String;
    potential_difference = 0.2u"V",
    pH_outer = 7.0,
    pMg_outer = 3.0,
    ionic_strength_outer = 0.25u"M",
    tolerance = 0.0,
)

    e_potential_difference =
        equilibrator.eq.Q_(string(uconvert(u"V", potential_difference)))

    outer_pH = equilibrator.eq.Q_(pH_outer)
    outer_pMg = equilibrator.eq.Q_(pMg_outer)
    outer_ionic_strength =
        equilibrator.eq.Q_(string(uconvert(u"M", float(ionic_strength_outer))))

    inner_rxn = equilibrator.cc.parse_reaction_formula(rxn_inner)
    outer_rxn = equilibrator.cc.parse_reaction_formula(rxn_outer)

    standard_dg_prime = equilibrator.cc.multicompartmental_standard_dg_prime(
        reaction_inner = inner_rxn,
        reaction_outer = outer_rxn,
        e_potential_difference = e_potential_difference,
        p_h_outer = outer_pH,
        ionic_strength_outer = outer_ionic_strength,
        p_mg_outer = outer_pMg,
        tolerance = tolerance,
    )

    return (
        standard_dg_prime.value.m_as("kJ/mol") ± standard_dg_prime.error.m_as("kJ/mol")
    )u"kJ/mol"
end


text:
module SMLMAssociationAnalysis_NCB

using Distributed
using LocalizationMicroscopy
using NearestNeighbors
using SharedArrays
using StatsPlots
using Statistics
using Tau

include("data/Molecule.jl")
include("data/ChannelData.jl")
include("data/Result.jl")
include("utilities.jl")
include("analysis/process.jl")
include("analysis/grouping.jl")
include("analysis/association.jl")
include("analysis/simulate.jl")
include("plots.jl")

export ChannelData, Molecule, Result, ResultSimulate
export ChannelDataOptimizing, ResultOptimizing
export getmolecules, exclusivenearestneighbors, montecarloaffinity
export moleculesplot, localizationsplot, insetplot, distanceprobabilityplot, neighborsplot, neighborsplot_forprint
export localizationsplot_forprint
export moleculesinsetplot, localizationsinsetplot, insetplot_forprint
export moleculesinsetplot_forprint, localizationsinsetplot_forprint
export randomcoordinates2d
export merge_close_molecules
export moleculesplot_sim
export simulate100, simulate0
end


text:
using Base.Iterators: partition
using ElasticArrays
using LinearAlgebra
using LyceumAI

# data buffer
struct DataBuffers{T<:AbstractFloat}
    phixudatabuf::ElasticArray{T}      # phi(x_t,u_t)
    xudatabuf::ElasticArray{T}         # x_t,u_t
    ydatabuf::ElasticArray{T}          # x_t+1
    xdatabuf::ElasticArray{T}          # x_t
    udatabuf::ElasticArray{T}          # u_t
    rewarddatabuf::ElasticArray{T}     # reward_t
    sizeofbuf::ElasticArray{Int}       # size of the buffer
    function DataBuffers{T}(ns::Int, no::Int, nu::Int, nf::Int) where T<:AbstractFloat
        new(ElasticArray{T}(undef,nf,0),
            ElasticArray{T}(undef,no+nu,0),
            ElasticArray{T}(undef,ns,0),
            ElasticArray{T}(undef,no,0),
            ElasticArray{T}(undef,nu,0),
            ElasticArray{T}(undef,1,0),
            ElasticArray{Int}(undef,1,0)
           )
    end
end

# LC3 structure
struct LC3{E, F, K, M}
    envsampler::E
    featurize::F
    PredictMat::K
    Hmax::Int      # maximum horizon
    N::Int         # number of trajectories
    ctrlrange::AbstractArray
    covW::AbstractArray
    Amat::AbstractArray
    covscale::AbstractFloat
    controller::M
    numfeatures::Int
    buffers::DataBuffers    # replay buffer sort of..
    function LC3(
                 env_tconstructor,
                 featurize,
                 PredictMat,
                 numfeatures,
                 ns,
                 no,
                 nu,
                 ctrlrange,
                 ridgereg,
                 covscale,
                 controller;

                 Hmax = 400,
                 N = 400,
                 prebuffer = nothing,
                )

        # check errors
        0 < Hmax <= N || throw(ArgumentError("Hmax must be in interval (0, N]"))
        0 < N || throw(ArgumentError("N must be > 0"))
        # find common type of PredictionMat
        #KT = eltype(PredictMat.W)
        #DT = LyceumAI.promote_modeltype(PredictMat)
        #if !isconcretetype(DT)
        #    DTnew = Shapes.default_datatype(DT)
        #    @warn "Could not infer model element type. Defaulting to $DTnew"
        #    DT = DTnew
        #end
        envsampler = env_tconstructor

        new{
            typeof(envsampler),
            typeof(featurize),
            typeof(PredictMat),
            typeof(controller)
           }(
             envsampler,
             featurize,
             PredictMat,
             Hmax,
             N,
             ctrlrange,
             Matrix(I(numfeatures)/ridgereg),
             zeros(numfeatures,no),
             covscale,
             controller,
             numfeatures,
             prebuffer == nothing ? DataBuffers{Float32}(ns,no,nu,numfeatures) : prebuffer,
            )
    end
end

#Get coherent data from multiple trajectory data
function getshifteddata(trajectory::AbstractArray)
    data1 = trajectory[:,1:end-1]
    data2 = trajectory[:,2:end]
    return (data1, data2)
end

#Clamp control inputs to the given range
function clampctrl!(actions::AbstractArray, ctrlrange::AbstractArray)
    nu = size(actions, 1)
    @inbounds for u=1:size(ctrlrange, 2)
        @simd for a=1:nu
            actions[u,a] = clamp(actions[u,a], ctrlrange[1,u], ctrlrange[2,u])
        end
    end
end

#Storing data and plot
function storedata!(batch::NamedTuple, buffers::DataBuffers, ctrlrange::AbstractArray, featurize)
    # Storing data into arrays : getshifteddata deals with multiple trajecotries
    xydata = getshifteddata(batch.observations)
    uvdata = getshifteddata(batch.actions)
    xdata  = first(xydata)
    udata  = first(uvdata)
    rewarddata = batch.rewards[1:end-1]
    clampctrl!(udata, ctrlrange)
    xudata = vcat(xdata, udata)
    PhiXUdata = featurize(xudata)
    ydata  = last(xydata) - xdata
    # save data into data buffers
    append!(buffers.phixudatabuf, PhiXUdata)
    append!(buffers.xudatabuf, xudata)
    append!(buffers.ydatabuf, ydata)
    append!(buffers.xdatabuf, xdata)
    append!(buffers.udatabuf, udata)
    append!(buffers.rewarddatabuf, rewarddata)
    append!(buffers.sizeofbuf, [size(xdata,2)])  # data size
    return xudata, ydata, PhiXUdata
end

#Iterator for PredictionMat Learning
function Base.iterate(featlearn::LC3{DT}, i = 1) where {DT}
    @unpack envsampler, featurize, PredictMat = featlearn
    @unpack Hmax, N, ctrlrange = featlearn
    @unpack covW, Amat, covscale, controller, numfeatures, buffers = featlearn

    if (i == 1) || (mod(i, 50) == 0); @info "Iterations:" i; end;

    #------------------Reset MPPI struct------------#
    reset!(controller)
    randreset!(envsampler)

    #-------------------Roll out--------------------#
    elapsed_sample = @elapsed begin
        opt = ControllerIterator((action, state, obs) -> getaction!(action, obs, controller),
                                 envsampler; T = Hmax, plotiter = 9999)
        for _ in opt # runs iterator
        end
        batch = opt.trajectory
    end

    #-------------------Store data------------------#
    newxu, newy, newPhiXU = storedata!(batch, buffers, ctrlrange, featurize)
    if (i == 1) || (mod(i, 50) == 0); @info "rewards" sum(batch.rewards); end;

    sizeofbuf = sum(buffers.sizeofbuf)

    #------------------Weight matrix----------------#
    # matrix inversion lemma / update covariance and compute mean model
    covW .= covW .- covW * newPhiXU * pinv(I+newPhiXU'*covW*newPhiXU) * newPhiXU' * covW
    Amat .+= newPhiXU * newy'
    meanW = covW * Amat

    prederr = norm(newy - meanW' * newPhiXU)/(norm(meanW' * newPhiXU) + 1e-6)
    #@info "Mean-model error" prederr

    # when covscale is set to zero, we use mean model
    if covscale == 0.
        # update weight matrix
        PredictMat.W .= meanW'
    # otherwise, we do Thompson sampling.  Each row is independent.
    else
        TestSampleMat = zeros(size(PredictMat.W,1),size(PredictMat.W,2))
        for row = 1:size(newy,1)
            samplerMat = MvNormal(meanW[:,row], Matrix(Hermitian(covW * covscale)))
            TestSampleMat[row,:] = rand(samplerMat)
        end
        # update weight matrix
        PredictMat.W .= TestSampleMat[:,:,1]
    end
    # update weight matrix in multiple learnedenvs for multi-threading
    for k in controller.envs
        k.PredictMat .= PredictMat.W
    end

    #------------------updating state---------------#
    result = (
              iter = i,
              elapsed_sampled = elapsed_sample,
              traj_reward = sum(batch.rewards),
              traj_length = mean(length, batch),
              prederr = prederr
             )

    return result, i + 1
end


text:
#=
P22 (*) Create a list containing all integers within a given range.
If first argument is smaller than second, produce a list in decreasing order.
Example:
* (range 4 9)
(4 5 6 7 8 9)
=#

function my_range(start, stop)
    function _my_range(start, coll)
        start > stop ? coll : _my_range(start + 1, vcat(coll, start))
    end
    _my_range(start, [])
end

@assert my_range(4, 9) == [4, 5, 6, 7, 8, 9]
@assert my_range(10, 10) == [10]

println("Tests passed: JL-22.jl")


text:
xs = [[0.,0.] [1.,1.] [2.,1.5]]
ts = [0.0, 1.0, 2.0]

pa = Path(xs,ts)

@test samplepath(pa, 0.3) == [0.3, 0.3] &&
      samplepath(pa, 1.3) == [1.3, 1.15]

# Seg1: [0.0, 1.0]
# Seg2: [1.0, 2.0]

@test norm(samplepath(pa, [0.3, 1.3, 1.4]) -
            [0.3 1.3 1.4; 0.3 1.15 1.2]) <= 1e-12

@test_throws AssertionError samplepath(pa, 10.)

xs = [[0., 0.] [1., 1.] [2., 1.5] [3., 5.0] [3.4, -2.0]]
ts = [0.0, 1.0, 2.0, 2.5, 2.7]

pb = Path(xs, ts)

N  = 10000
ss = range(0.0, stop=ts[end], length=N)

samples = samplepath(pb, ss)
@test norm(sum(samples, dims=2) / N - pathmean(pb)) / norm(pathmean(pb)) <= 5e-3

#####
Random.seed!(15)
N  = 500
xs = zeros(2, N)
ts = zeros(N)
xs[:, 1] = randn(2)
for s ∈ 2:N
    xs[:, s] = randn(2)
    ts[s] = ts[s-1] + rand()
end
pb = Path(xs, ts)
N  = 10000
ss = range(0.0, stop=ts[end], length=N)
samples = samplepath(pb, ss)
@test norm(sum(samples, dims=2) / N - pathmean(pb)) / norm(pathmean(pb)) <= 5e-3


text:
# Index of representation element of type ElemT
"""
    Index{N, T, ElemT}

Index of an element of type `ElemT` in a `Rep{N, T}`.
"""
struct Index{N, T, ElemT}
    value::Int
end

# Type of the value associated to this index
valuetype(::Union{Index{N, T, ElemT}, Type{Index{N, T, ElemT}}}) where {N, T, ElemT} = ElemT

const HyperPlaneIndex{N, T} = Index{N, T, <:HyperPlane{N, T}}
const HalfSpaceIndex{N, T} = Index{N, T, <:HalfSpace{N, T}}
const HIndex{N, T} = Union{HyperPlaneIndex{N, T}, HalfSpaceIndex{N, T}}

#const SymPointIndex{N, T} = Index{N, T, <:SymPoint{N, T}}
const PointIndex{N, T} = Index{N, T, <:AbstractPoint{N, T}}
#const PIndex{N, T} = Union{SymPointIndex{N, T}, PointIndex{N, T}}
const PIndex{N, T} = PointIndex{N, T}
const LineIndex{N, T} = Index{N, T, <:Line{N, T}}
const RayIndex{N, T} = Index{N, T, <:Ray{N, T}}
const RIndex{N, T} = Union{LineIndex{N, T}, RayIndex{N, T}}
const VIndex{N, T} = Union{PIndex{N, T}, RIndex{N, T}}
islin(::Union{Index{N, T, ElemT}, Type{Index{N, T, ElemT}}}) where {N, T, ElemT} = islin(ElemT)
ispoint(::Union{Index{N, T, ElemT}, Type{Index{N, T, ElemT}}}) where {N, T, ElemT} = ispoint(ElemT)

"""
    Indices{N, T, ElemT, RepT<:Rep{N, T}}

Iterator over the indices of the elements of type `ElemT` of the field `rep`.
"""
struct Indices{N, T, ElemT, RepT<:Rep{N, T}}
    rep::RepT
    function Indices{N, T, ElemT}(rep) where {N, T, ElemT}
        new{N, T, ElemT, typeof(rep)}(rep)
    end
end

Base.eltype(::Indices{N, T, ElemT}) where {N, T, ElemT} = Index{N, T, ElemT}
valuetype(idxs::Indices) = valuetype(eltype(idxs))

const HyperPlaneIndices{N, T, RepT} = Indices{N, T, <:HyperPlane{N, T}, RepT}
const HalfSpaceIndices{N, T, RepT} = Indices{N, T, <:HalfSpace{N, T}, RepT}
const HIndices{N, T, RepT} = Union{HyperPlaneIndices{N, T, RepT}, HalfSpaceIndices{N, T, RepT}}

#const SymPointIndices{N, T, RepT} = Indices{N, T, <:SymPoint{N, T}, RepT}
const PointIndices{N, T, RepT} = Indices{N, T, <:AbstractPoint{N, T}, RepT}
#const PIndices{N, T, RepT} = Union{SymPointIndices{N, T, RepT}, PointIndices{N, T, RepT}}
const PIndices{N, T, RepT} = PointIndices{N, T, RepT}
const LineIndices{N, T, RepT} = Indices{N, T, <:Line{N, T}, RepT}
const RayIndices{N, T, RepT} = Indices{N, T, <:Ray{N, T}, RepT}
const RIndices{N, T, RepT} = Union{LineIndices{N, T, RepT}, RayIndices{N, T, RepT}}
const VIndices{N, T, RepT} = Union{PIndices{N, T, RepT}, RIndices{N, T, RepT}}

function Base.next(idxs::Indices{N, T, ElemT}, idx::Index{N, T, ElemT}) where {N, T, ElemT}
    nextidx = nextindex(idxs.rep, idx)
    idx, nextidx
end

repfor(p, ::Type{<:HRepElement}) = hrep(p)
repfor(p, ::Type{<:VRepElement}) = vrep(p)
Base.length(idxs::Indices{N, T, ElemT, <:Polyhedron{N, T}}) where {N, T, ElemT} = length(Indices{N, T, ElemT}(repfor(idxs.rep, ElemT)))
Base.isempty(idxs::Indices{N, T, ElemT, <:Polyhedron{N, T}}) where {N, T, ElemT} = isempty(Indices{N, T, ElemT}(repfor(idxs.rep, ElemT)))
Base.start(idxs::Indices{N, T, ElemT, <:Polyhedron{N, T}}) where {N, T, ElemT} = start(Indices{N, T, ElemT}(repfor(idxs.rep, ElemT)))
Base.done(idxs::Indices{N, T, ElemT, <:Polyhedron{N, T}}, idx::Index{N, T, ElemT}) where {N, T, ElemT} = done(Indices{N, T, ElemT}(repfor(idxs.rep, ElemT)), idx)
Base.get(p::Polyhedron{N, T}, idx::Index{N, T, ElemT}) where {N, T, ElemT} = get(repfor(p, ElemT), idx)
nextindex(p::Polyhedron{N, T}, idx::Index{N, T, ElemT}) where {N, T, ElemT} = nextindex(repfor(p, ElemT), idx)

"""
The representation `rep` does not contain any `elem`.
"""
macro norepelem(rep, elem)
    idxs = Symbol(string(elem) * "Indices")
    idx = Symbol(string(elem) * "Index")
    quote
        Base.length(idxs::$idxs{N, T, <:$rep{N, T}}) where {N, T} = 0
        Base.isempty(idxs::$idxs{N, T, <:$rep{N, T}}) where {N, T} = true
        Base.start(idxs::$idxs{N, T, <:$rep{N, T}}) where {N, T} = eltype(idxs)(0)
        Base.done(idxs::$idxs{N, T, <:$rep{N, T}}, ::$idx{N, T}) where {N, T} = true
    end
end

abstract type HAffineSpace{N, T} <: HRepresentation{N, T} end
@norepelem HAffineSpace HalfSpace

_promote_reptype(P1::Type{<:HAffineSpace}, ::Type{<:HAffineSpace}) = P1
_promote_reptype(P1::Type{<:HAffineSpace}, ::Type{<:HRep}) = hreptype(P1)

abstract type VPolytope{N, T} <: VRepresentation{N, T} end
@norepelem VPolytope Line
@norepelem VPolytope Ray

_promote_reptype(P1::Type{<:VPolytope}, ::Type{<:VPolytope}) = P1
_promote_reptype(P1::Type{<:VPolytope}, ::Type{<:VRep}) = vreptype(P1)

abstract type VSymPolytope{N, T} <: VPolytope{N, T} end
@norepelem VSymPolytope Point

abstract type VCone{N, T} <: VRepresentation{N, T} end
#@norepelem VCone SymPoint
# See issue #28
Base.length(idxs::PointIndices{N, T, <:VCone{N, T}}) where {N, T} = hasallrays(idxs.rep) ? 1 : 0
Base.isempty(idxs::PointIndices{N, T, <:VCone{N, T}}) where {N, T} = !hasallrays(idxs.rep)
Base.start(idxs::PointIndices{N, T, <:VCone{N, T}}) where {N, T} = eltype(idxs)(hasallrays(idxs.rep) ? 1 : 2)
Base.done(::PointIndices{N, T, <:VCone{N, T}}, idx::PointIndex{N, T}) where {N, T} = idx.value > 1
Base.get(L::VCone{N, T}, ::PointIndex{N, T}) where {N, T} = origin(vvectortype(typeof(L)), FullDim{N}())
nextindex(::VCone{N, T}, idx::PointIndex{N, T}) where {N, T} = typeof(idx)(idx.value + 1)

_promote_reptype(P1::Type{<:VCone}, ::Type{<:VCone}) = P1
_promote_reptype(P1::Type{<:VCone}, ::Type{<:VRep}) = vreptype(P1)

abstract type VLinearSpace{N, T} <: VCone{N, T} end
@norepelem VLinearSpace Ray

_promote_reptype(P1::Type{<:VLinearSpace}, ::Type{<:VLinearSpace}) = P1
_promote_reptype(P1::Type{<:VLinearSpace}, ::Type{<:VCone}) = conetype(P1)
_promote_reptype(P1::Type{<:VLinearSpace}, ::Type{<:VRep}) = vreptype(P1)

"""
The representation `rep` contain the elements `elem` inside a vector in the field `field`.
"""
macro vecrepelem(rep, elem, field)
    idxs = Symbol(string(elem) * "Indices")
    idx = Symbol(string(elem) * "Index")
    esc(quote
        Base.length(idxs::$idxs{N, T, <:$rep{N, T}}) where {N, T} = length(idxs.rep.$field)
        Base.isempty(idxs::$idxs{N, T, <:$rep{N, T}}) where {N, T} = isempty(idxs.rep.$field)
        Base.start(idxs::$idxs{N, T, <:$rep{N, T}}) where {N, T} = eltype(idxs)(1)
        Base.done(idxs::$idxs{N, T, <:$rep{N, T}}, idx::$idx{N, T}) where {N, T} = idx.value > length(idxs)
        Base.get(rep::$rep{N, T}, idx::$idx{N, T}) where {N, T} = rep.$field[idx.value]
        nextindex(::$rep{N, T}, idx::$idx{N, T}) where {N, T} = typeof(idx)(idx.value + 1)
    end)
end

"""
The representation `rep` contain the elements `elem` inside a representation in the field `field`.
"""
macro subrepelem(rep, elem, field)
    idxst = Symbol(string(elem) * "Indices")
    idxs = :(Polyhedra.$idxst)
    idxt = Symbol(string(elem) * "Index")
    idx = :(Polyhedra.$idxt)
    subidxs = :(Polyhedra.Indices{N, T, Polyhedra.valuetype(idxs)}(idxs.rep.$field))
    esc(quote
        Base.length(idxs::$idxs{N, T, <:$rep{N, T}}) where {N, T} = length($subidxs)
        Base.isempty(idxs::$idxs{N, T, <:$rep{N, T}}) where {N, T} = isempty($subidxs)
        Base.start(idxs::$idxs{N, T, <:$rep{N, T}}) where {N, T} = start($subidxs)
        Base.done(idxs::$idxs{N, T, <:$rep{N, T}}, idx::$idx{N, T}) where {N, T} = done($subidxs, idx)
        Base.get(rep::$rep{N, T}, idx::$idx{N, T}) where {N, T} = get(rep.$field, idx)
        Polyhedra.nextindex(rep::$rep{N, T}, idx::$idx{N, T}) where {N, T} = Polyhedra.nextindex(rep.$field, idx)
    end)
end


text:
##############################################################################
##
## TODO: dtype == :complex
##       See minFunc code for examples of how this could be done
##
##############################################################################

##############################################################################
##
## Graveyard of Alternative Functions
##
## In Nodedal and Wright, epsilon was calculated as
##
##   epsilon = sqrt(eps())
##
## In the examples I've tested, this strategy is best when x is small
##  and becomes very bad when x is larger
##
## In previous iteration of derivative_numer, epsilon was calculated as
##
##   epsilon = eps(max(1.0, abs(x)))^(1/3)
##
## In the examples I've tested, this strategy is bad when x is small
##  and becomes better when x is larger
##
## The current approach comes from minFunc
##
##    epsilon = 2 * sqrt(1e-12) * (1 + norm(x))
##
## TODO: Perform a systematic comparison across several functions
##       and several ranges of inputs
##
##############################################################################

##############################################################################
##
## Derivative of f: R -> R
##
##############################################################################

function finite_difference{T <: Number}(f::Function,
                                        x::T,
                                        dtype::Symbol)
    epsilon = 2 * sqrt(1e-12) * (1 + norm(x))

    if dtype == :forward
        return (f(x + epsilon) - f(x)) / epsilon
    elseif dtype == :central
        return (f(x + epsilon) - f(x - epsilon)) / (2 * epsilon)
    else
        error("dtype must be :forward or :central")
    end
end
finite_difference(f::Function, x::Number) = finite_difference(f, x, :central)

##############################################################################
##
## Gradient of f: R^n -> R
##
##############################################################################

function finite_difference{T <: Number}(f::Function,
                                        x::Vector{T},
                                        dtype::Symbol)
    epsilon = 2 * sqrt(1e-12) * (1 + norm(x))

    # What is the dimension of x?
    n = length(x)

    # Storage for forward differences
    differential = Array(Float64, n)

    # Iterate over each dimension of the gradient separately.
    # Use xplusdx to store x + dx instead of creating a new vector on each pass.
    # Use xminusdx to store x - dx instead of creating a new vector on each pass.
    if dtype == :forward
        # Establish a baseline value of f(x).
        f_x = f(x)
        xplusdx = copy(x)
        for i = 1:n
            xplusdx[i] = x[i] + epsilon
            differential[i] = (f(xplusdx) - f_x) / epsilon
            xplusdx[i] = x[i]
        end
    elseif dtype == :central
        xplusdx, xminusdx = copy(x), copy(x)
        for i = 1:n
            xplusdx[i], xminusdx[i] = x[i] + epsilon, x[i] - epsilon
            differential[i] = (f(xplusdx) - f(xminusdx)) / (2 * epsilon)
            xplusdx[i], xminusdx[i] = x[i], x[i]
        end
    else
        error("dtype must be :forward or :central")
    end

    # Return the estimated gradient.
    return differential
end
function finite_difference{T <: Number}(f::Function, x::Vector{T})
    finite_difference(f, x, :central)
end

##############################################################################
##
## Jacobian derivative of f: R^n -> R^m
##
##############################################################################

function finite_difference_jacobian{T <: Number}(f::Function, x::Vector{T}, dtype::Symbol)
    epsilon = 2 * sqrt(1e-12) * (1 + norm(x))

    # What is the dimension of x?
    n = length(x)

    # Establish a baseline for f_x
    f_x = f(x)

    # Initialize the Jacobian matrix
    J = zeros(length(f_x), n)

    # Iterate over each dimension of the gradient separately.
    if dtype == :forward
        xplusdx = copy(x)
        for i = 1:n
            xplusdx[i] = x[i] + epsilon
            J[:, i] = (f(xplusdx) - f_x) / epsilon
            xplusdx[i] = x[i]
        end
        return J
    elseif dtype == :central
        xplusdx, xminusdx = copy(x), copy(x)
        for i = 1:n
            xplusdx[i], xminusdx[i] = x[i] + epsilon, x[i] - epsilon
            J[:, i] = (f(xplusdx) - f(xminusdx)) / (2 * epsilon)
            xplusdx[i], xminusdx[i] = x[i], x[i]
        end
        return J
    else
        error("dtype must :forward or :central")
    end
end

##############################################################################
##
## Second derivative of f: R^n -> R
##
##############################################################################

function finite_difference_hessian(f::Function, g::Function, x::Number, dtype::Symbol)
    epsilon = 2 * sqrt(1e-12) * (1 + norm(x))

    if dtype == :forward
        return (g(x + epsilon) - g(x)) / epsilon
    elseif dtype == :central
        return (g(x + epsilon) - g(x - epsilon)) / (2 * epsilon)
    else
        error("dtype must :forward or :central")
    end
end
function finite_difference_hessian(f::Function, x::Number, dtype::Symbol)
    finite_difference_hessian(f, derivative(f), x, dtype)
end
function finite_difference_hessian(f::Function, x::Number)
    finite_difference_hessian(f, derivative(f), x, :central)
end

##############################################################################
##
## Hessian of f: R^n -> R
##
##############################################################################

function finite_difference_hessian{T <: Number}(f::Function, g::Function, x::Vector{T}, dtype::Symbol)
    epsilon = 2 * sqrt(1e-12) * (1 + norm(x))

    # What is the dimension of x?
    n = length(x)

    # Initialize an empty Hessian
    H = Array(Float64, n, n)

    if dtype == :forward
        g_x = g(x)
        xplusdx = copy(x)
        for i in 1:n
            xplusdx[i] = x[i] + epsilon
            H[:, i] = (g(xplusdx) - g_x) / epsilon
            xplusdx[i] = x[i]
        end
        symmetrize!(H)
        return H
    elseif dtype == :central
        xplusdx, xminusdx = copy(x), copy(x)
        for i = 1:n
            xplusdx[i], xminusdx[i] = x[i] + epsilon, x[i] - epsilon
            H[:, i] = (g(xplusdx) - g(xminusdx)) / (2 * epsilon)
            xplusdx[i], xminusdx[i] = x[i], x[i]
        end
        symmetrize!(H)
        return H
    else
        error("dtype must :forward or :central")
    end
end
function finite_difference_hessian{T <: Number}(f::Function, x::Vector{T}, dtype::Symbol)
    finite_difference_hessian(f, gradient(f), x, dtype)
end
function finite_difference_hessian{T <: Number}(f::Function, x::Vector{T})
    finite_difference_hessian(f, gradient(f), x, :central)
end

##############################################################################
##
## Taylor Series based estimates of first and second derivatives
##
## TODO: Fill in multivariate equivalents
##
##############################################################################

# Higher precise finite difference method based on Taylor series approximation.
# h is the stepsize
function taylor_finite_difference(f::Function, x::Float64, h::Float64, dtype::Symbol)
    if dtype == :forward
        f_x = f(x)
        d = 2^3 * (2^2 * (f(x + h) - f_x) - (f(x + 2 * h) - f_x))
        d += - (2^2 * (f(x + 2 * h) - f_x) - (f(x + 4 * h) - f_x))
        d /= 3 * 2^2 * h
    elseif dtype == :central
        d = 4^5 * (2^3 * (f(x + h) - f(x - h)) - (f(x + 2 * h) - f(x - 2 * h)))
        d -= 2^3 * (f(x + 4 * h) - f(x - 4 * h)) - (f(x + 8 * h) - f(x - 8 * h))
        d /= (4^5 * (2^4 - 2^2) - (2^6 - 2^4)) * h
    else
        error("dtype must be :forward or :central")
    end
    return d
end
function taylor_finite_difference(f::Function, x::Float64, h::Float64)
    taylor_finite_difference(f, x, h, :central)
end
function taylor_finite_difference(f::Function, x::Float64, dtype::Symbol)
    taylor_finite_difference(f, x, 10e-4, dtype)
end
function taylor_finite_difference(f::Function, x::Float64)
    taylor_finite_difference(f, x, 10e-4)
end

function taylor_finite_difference_hessian(f::Function, x::Float64, h::Float64)
    f_x = f(x)
    d = 4^6 * (2^4 * (f(x + h) + f(x - h) - 2 * f_x) - (f(x + 2 * h) + f(x - 2 * h) - 2 * f_x))
    d += - (2^4 * (f(x + 4 * h) + f(x - 4 * h) - 2 * f_x) - (f(x + 8 * h) + f(x - 8 * h) - 2 * f_x))
    return d / (3 * 2^6 * (2^8 - 1) * h^2)
end

##############################################################################
##
## TODO: Implement directional_derivative()
##
##############################################################################

# The function "dirderivative" calculates directional derivatives in the direction v.
# The function supplied must have the form Vector{Float64} -> Float64
# function dirderivative(f::Function, v::Vector{Float64}, x0::Vector{Float64}, h::Float64, twoside::Bool)
#     derivative(t::Float64 -> f(x0 + v*t) / norm(v), 0.0, h, twoside)
# end
# function dirderivative(f::Function, v::Vector{Float64}, x0::Vector{Float64}, h::Float64)
#     dirderivative(f, v, x0, h, true)
# end
# function dirderivative(f::Function, v::Vector{Float64}, x0::Vector{Float64}, )
#     derivative(f, v, x0, 0.0001)
# end
# function dirderivative(f::Function, v::Vector{Float64})
#     x -> dirderivative(f, v, x)
# end


text:
# dual only versus dual only no scaling

using JLD

include("perfplots.jl")

timeshift = 10.0

# produced by process_csv.jl in perfprofile mode
jldsource = joinpath(dirname(@__FILE__),"subpsep_perf.jld")

d = load(jldsource)
time_table = d["time_table"] + timeshift
itercount_table = d["itercount_table"]
solvers = d["solvers"]
@assert solvers == ["PAJ_CPLEX_MOSEK_dualonly","PAJ_CPLEX_MOSEK_primonlynorelax", "PAJ_CPLEX_MOSEK_msddualonly", "PAJ_CPLEX_MOSEK_msdprimonlynorelax"]
instance_names = d["instance_names"]

Plots.scalefontsizes(2.1)

performance_profile(time_table, ["iter. subp","iter. sep","MSD subp","MSD sep"],logscale=false, ymax=1.0, xmax=15, linewidth=3)
Plots.savefig(joinpath(dirname(@__FILE__),"septime.eps"))

max_stars_repo_name: mlubin/PajaritoSupplement

max_issues_repo_name: mlubin/PajaritoSupplement

max_forks_repo_name: mlubin/PajaritoSupplement


text:
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# FILE CONTENTS:
#     SUMMARY: Integration schemes
#
#     INTEGRATION METHODS
#         midpoint: midpoint or trapezoidal integration
#         rk4: Runge-Kutta 4
#         rk3: Runge-Kutta 3
#     OTHER METHODS
#         f_augmented!: Create function with augmented state and control input
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

# TODO: Change S to Z

default_integration() = rk4

"""
$(SIGNATURES)
In place Midpoint integration
Defines methods for both separated and augmented forms. Returns a discrete version
of a continuous dynamics function.
# Arguments
* f!: in place dynamics function, i.e. `f!(x?,x,u)`
* dt: time step
"""
function midpoint(f!::Function, dt::Float64)
   fd!(xdot,x,u,dt=dt) = begin
       f!(xdot,x,u)
       xdot .*= dt/2.
       f!(xdot, x + xdot, u)
       copyto!(xdot,x + xdot*dt)
   end
end

function midpoint_uncertain(f!::Function, dt::Float64)
   fd!(xdot,x,u,w,dt=dt) = begin
       f!(xdot,x,u,w)
       xdot .*= dt/2.
       f!(xdot, x + xdot, u, w)
       copyto!(xdot,x + xdot*dt)
   end
end

function midpoint_implicit(f::Function,n::Int,m::Int,dt::T) where T
   # get estimate of X[k+1] from explicit midpoint
   f_aug(ẋ,z) = f(ẋ,z[1:n],z[n .+ (1:m)])
   ∇f(x,u) = ForwardDiff.jacobian(f_aug,zero(x),[x;u])

   fd(y,x,u,dt=dt) = begin
       fc = zero(x)

       copyto!(y, x)
       g = Inf
       cnt = 0
       while norm(g) > 1.0e-12
           cnt += 1
           if cnt > 1000
               error("Integration convergence fail")
           end

           Xm = 0.5*(x + y)
           f(fc,Xm,u)
           g = y - x - dt*fc

           A = ∇f(Xm,u)[:,1:n]

           ∇g = Diagonal(I,n) - 0.5*dt*A
           δy = -∇g\g

           y .+= δy
       end
   end
end

function midpoint_implicit_uncertain(f::Function,n::Int,m::Int,r::Int,dt::T) where T
   # get estimate of X[k+1] from explicit midpoint
   f_aug(ẋ,z) = f(ẋ,z[1:n],z[n .+ (1:m)],z[(n+m) .+ (1:r)])
   ∇f(x,u,w) = ForwardDiff.jacobian(f_aug,zeros(eltype(x),n),[x;u;w])
   fd(y,x,u,w,dt=dt) = begin
       fc = zero(x)

       copyto!(y, x )

       # iterate to solve implicit midpoint step
       g = Inf
       cnt = 0
       while norm(g) > 1.0e-12
           cnt += 1
           if cnt > 1000
               error("Integration tolerance failed")
           end
           Xm = 0.5*(x + y)
           f(fc,Xm,u,w)
           g = y - x - dt*fc

           A = ∇f(Xm,u,w)[:,1:n]

           ∇g = Diagonal(I,n) - 0.5*dt*A
           δy = -∇g\g

           y .+= δy
       end
   end
end

"""
$(SIGNATURES)
In place Runge Kutta 4 integration
Defines methods for both separated and augmented forms. Returns a discrete version
of a continuous dynamics function.
# Arguments
* f!: in place dynamics function, i.e. `f!(x?,x,u)` for separate or `f!(S?,S)` for augmented dynamics
* dt: time step
"""
function rk4(f!::Function, dt::Float64)
   # Runge-Kutta 4
   fd!(xdot,x,u,dt=dt) = begin
       k1 = k2 = k3 = k4 = zero(xdot)
       f!(k1, x, u);         k1 *= dt;
       f!(k2, x + k1/2, u); k2 *= dt;
       f!(k3, x + k2/2, u); k3 *= dt;
       f!(k4, x + k3, u);    k4 *= dt;
       copyto!(xdot, x + (k1 + 2*k2 + 2*k3 + k4)/6)
   end
end

function rk4_uncertain(f!::Function, dt::Float64)
   # Runge-Kutta 4
   fd!(xdot,x,u,w,dt=dt) = begin
       k1 = k2 = k3 = k4 = zero(xdot)
       f!(k1, x, u, w);         k1 *= dt;
       f!(k2, x + k1/2, u, w); k2 *= dt;
       f!(k3, x + k2/2, u, w); k3 *= dt;
       f!(k4, x + k3, u, w);    k4 *= dt;
       copyto!(xdot, x + (k1 + 2*k2 + 2*k3 + k4)/6)
   end
end

"""
$(SIGNATURES)
In place Runge Kutta 3 integration
Defines methods for both separated and augmented forms. Returns a discrete version
of a continuous dynamics function.
# Arguments
* f!: in place dynamics function, i.e. `f!(x?,x,u)` for separate or `f!(S?,S)` for augmented dynamics
* dt: time step
"""
function rk3(f!::Function, dt::Float64)
       # Runge-Kutta 3 (zero order hold)
   fd!(xdot,x,u,dt=dt) = begin
       k1 = k2 = k3 = zero(x)
       f!(k1, x, u);               k1 *= dt;
       f!(k2, x + k1/2, u);       k2 *= dt;
       f!(k3, x - k1 + 2*k2, u);  k3 *= dt;
       copyto!(xdot, x + (k1 + 4*k2 + k3)/6)
   end
end

function rk3_uncertain(f!::Function, dt::Float64)
       # Runge-Kutta 3 (zero order hold)
   fd!(xdot,x,u,w,dt=dt) = begin
       k1 = k2 = k3 = zero(x)
       f!(k1, x, u, w);               k1 *= dt;
       f!(k2, x + k1/2, u, w);       k2 *= dt;
       f!(k3, x - k1 + 2*k2, u, w);  k3 *= dt;
       copyto!(xdot, x + (k1 + 4*k2 + k3)/6)
   end
end

function rk3_implicit(f::Function,n::Int,m::Int,dt::T) where T
   f_aug(ẋ,z) = f(ẋ,z[1:n],z[n .+ (1:m)])
   ∇f(x,u) = ForwardDiff.jacobian(f_aug,zero(x),[x;u])

   fd!(y,x,u,dt=dt) = begin
       fc1 = fc2 = fc3 = zero(x)

       g = Inf
       cnt = 0
       copyto!(y,x)
       while norm(g) > 1.0e-12
           cnt += 1
           if cnt > 1000
               println("norm: \n $(norm(g))")
               error("Integration convergence failed")
           end
           f(fc1,x,u)
           f(fc3,y,u)

           Xm = 0.5*(x + y) + dt/8*(fc1 - fc3)
           f(fc2,Xm,u)

           g = y - x - dt/6*fc1 - 4/6*dt*fc2 - dt/6*fc3

           A1 = ∇f(Xm,u)[:,1:n]
           A2 = ∇f(y,u)[:,1:n]

           ∇g = Diagonal(I,n) - 4/6*dt*A1*(0.5*Diagonal(I,n) - dt/8*A2) - dt/6*A2
           δy = -∇g\g

           y .+= δy
       end
   end
end

function rk3_implicit_uncertain(f::Function,n::Int,m::Int,r::Int,dt::T) where T
   f_aug(ẋ,z) = f(ẋ,z[1:n],z[n .+ (1:m)],z[(n+m) .+ (1:r)])
   ∇f(x,u,w) = ForwardDiff.jacobian(f_aug,zero(x),[x;u;w])

   fd!(y,x,u,w,dt=dt) = begin
       fc1 = fc2 = fc3 = zero(x)

       g = Inf
       cnt = 0
       copyto!(y,x)

       while norm(g) > 1.0e-12
           cnt += 1
           if cnt > 1000
               println(norm(g))
               error("Integration convergence failed")
           end
           f(fc1,x,u,w)
           f(fc3,y,u,w)

           Xm = 0.5*(x + y) + dt/8*(fc1 - fc3)
           f(fc2,Xm,u,w)

           g = y - x - dt/6*fc1 - 4/6*dt*fc2 - dt/6*fc3

           A1 = ∇f(Xm,u,w)[:,1:n]
           A2 = ∇f(y,u,w)[:,1:n]

           ∇g = Diagonal(I,n) - 4/6*dt*A1*(0.5*Diagonal(I,n) - dt/8*A2) - dt/6*A2
           δy = -∇g\g

           y .+= δy
           # copyto!(y,y+δy)
       end
   end
end

## DifferentialEquations.jl
# NOTE: not fast
function DiffEqIntegrator(f!::Function, dt::Float64, integrator::Symbol, n::Int, m::Int)
   function f_aug(z,p,t)
       ż = zero(z)
       f!(view(ż,1:n),z[1:n],z[n .+ (1:m)])
       ż
   end

   function fd_ode(y,x,u,dt=dt)
       _tf = dt
       _t0 = 0.

       u0=vec([x;u])
       tspan = (_t0,_tf)
       pro = ODEProblem(f_aug,u0,tspan)
       sol = OrdinaryDiffEq.solve(pro,eval(integrator)(),dt=dt,save_everystep=false)#,verbose=false)
       copyto!(y,sol.u[end][1:n])
   end
end

function DiffEqIntegratorUncertain(f!::Function, dt::Float64, integrator::Symbol, n::Int, m::Int, r::Int)
   function f_aug(z,p,t)
       ż = zero(z)
       f!(view(ż,1:n),z[1:n],z[n .+ (1:m)],z[(n+m) .+ (1:r)])
       ż
   end

   function fd_ode(y,x,u,w,dt=dt)
       _tf = dt
       _t0 = 0.
       u0=vec([x;u;w])
       tspan = (_t0,_tf)
       pro = ODEProblem(f_aug,u0,tspan)
       sol = OrdinaryDiffEq.solve(pro,eval(integrator)(),dt=dt,save_everystep=false)#,verbose=false)
       copyto!(y,sol.u[end][1:n])
   end
end

"""
$(SIGNATURES)
Converts a separated dynamics function into an augmented dynamics function
"""
function f_augmented!(f!::Function, n::Int, m::Int)
   f_aug!(dS::AbstractArray, S::Array) = f!(dS, S[1:n], S[n+1:n+m])
end


text:
using Printf, LBFGSB

const START = b"START"
const FG = b"FG"
const STOP = b"STOP"
const NEW_X = b"NEW_X"

const STOP2 = b"ST"
const NEW_X2 = b"NE"

function copyval!(x, copytoinds, copyfromind)
    for i in copytoinds
        @inbounds x[i] = x[copyfromind]
    end
end

function calllbfgsb!(
    ncm,
    G,
    H,
    tol,
    t,
    L,
    τ,
    α,
    σ;
    method::Symbol = :IAPG,
    innertol::Float64 = 1e-2,
    maxfgcalls::Int64 = 100,
    lbfgsbprintlevel::Int64 = -1,
    exact::Bool = false,
    verbose::Bool = false,
    cleanvals::Bool = true,
    scaleX::Bool = true,
)

    n = ncm.n
    g = ncm.g

    V = ncm.V
    Z = ncm.Z

    wa = ncm.wa
    iwa = ncm.iwa
    nbd = ncm.nbd
    lower = ncm.lower
    upper = ncm.upper
    task = ncm.task
    task2 = ncm.task2
    csave = ncm.csave
    lsave = ncm.lsave
    isave = ncm.isave
    dsave = ncm.dsave
    nRef = ncm.nRef
    mRef = ncm.mRef
    iprint = ncm.iprint
    fRef = ncm.fRef
    factr = ncm.factr
    pgtol = ncm.pgtol

    iprint[] = lbfgsbprintlevel

    proj = ncm.proj
    res = ncm.res

    Xnew = res.X
    y = res.y
    Λ = res.Λ

    fvals = res.fvals
    resvals = res.resvals
    distvals = res.distvals
    fgcountRef = res.fgcountRef
    εRef = res.εRef

    # Reset L-BFGS-B arrays
    fill!(wa, 0.0)
    fill!(iwa, 0)
    fill!(task, Cuchar(' '))
    fill!(csave, Cuchar(' '))
    fill!(lsave, 0)
    fill!(isave, 0)
    fill!(dsave, 0.0)

    fgcalls = 0
    linesearchcount = 0

    StopBFGS = false
    successful = true
    computeV = (method == :IAPG || method == :IR || method == :IER)

    # "We start the iteration by initializing task."
    copyto!(task, START)

    while !StopBFGS

        # This is the call to the L-BFGS-B code.
        setulb(
            nRef,
            mRef,
            y,
            lower,
            upper,
            nbd,
            fRef,
            g,
            factr,
            pgtol,
            wa,
            iwa,
            task,
            iprint,
            csave,
            lsave,
            isave,
            dsave,
        )

        if cleanvals && linesearchcount > 1
            a = fgcountRef[] - linesearchcount + 1
            b = fgcountRef[]
            copytoinds = a:b-1
            copyval!(fvals, copytoinds, b)
            copyval!(resvals, copytoinds, b)
            copyval!(distvals, copytoinds, b)
        end

        if task2 == FG
            if fgcalls >= maxfgcalls
                copyto!(task, STOP)
            else
                fRef[] = dualobj!(
                    ncm,
                    G,
                    H,
                    L,
                    τ,
                    α,
                    method,
                    computeV = computeV,
                    scaleX = scaleX,
                )
                fgcalls += 1
                if fgcalls > 1
                    linesearchcount += 1
                end

                if resvals[fgcountRef[]] < tol
                    copyto!(task, STOP)
                elseif !exact
                    εRef[] = symdot(Xnew, Λ)
                    ε = max(0.0, εRef[])
                    if method == :IAPG
                        δ = fronorm(V, proj.work)
                        lhs = 1/√L * δ
                        rhs = innertol / (√2 * t)
                        condition = (lhs ≤ rhs) && (ε ≤ innertol / (2*t^2))
                    else
                        δ = fronorm(V, proj.work)
                        dist = distvals[fgcountRef[]]
                        if method == :IR
                            lhs = (τ * δ)^2 + 2τ * ε * L
                            rhs = L * ((1 - τ) * L - α * τ) * dist^2
                            condition = (lhs ≤ rhs)
                        elseif method == :IER
                            lhs = (α * δ)^2 + 2α * ε
                            rhs = (σ * dist)^2
                            condition = (lhs ≤ rhs)
                        end
                    end

                    if condition
                        copyto!(task, STOP)
                    end
                end
            end

        elseif task2 == NEW_X2
            linesearchcount = 0

        else
            StopBFGS = true
            if !exact && task2 != STOP2
                verbose && @printf("\ntask = %s\n", String(copy(task)))
                successful = false
            end
        end

    end

    return successful
end

max_stars_repo_name: krislock/IR-FISTA

max_issues_repo_name: krislock/IR-FISTA

max_forks_repo_name: krislock/IR-FISTA


text:
include("1.jl")
function standart(r::Robot,n::Int64)
    flag1 = 1
    while (!isborder(r,Nord))
        y = n
        while (y > 0)
            putmarkersus(r,flag1,n,Ost)
            move!(r,Nord)
            y -= 1
        end
        flag1 = mod(flag1 + 1, 2)
    end
    if (mod(flag1 + 1, 2) == 1)
        putmarkersus(r,flag1,n,Ost)
    end
end

function putmarkersus(r::Robot,flag::Int64,n::Int64,side::HorizonSide)
    flag0 = flag
    while (!isborder(r,side))
        y = n
        while (y > 0)
            if ((flag0 == 0) && (!isborder(r,side)))
                move!(r,side)
            end
            if ((flag0 == 1) && (!isborder(r,side)))
                putmarker!(r)
                move!(r,side)
            end
            y -= 1
        end
        flag0 = mod(flag0 + 1, 2)
    end
    if (mod(flag0 + 1, 2) == 1)
        putmarker!(r)
    end
    while (!isborder(r,inverse(side)))
        move!(r,inverse(side))
    end
end

max_stars_repo_name: bubasooo/proggggggggg

max_issues_repo_name: bubasooo/proggggggggg

max_forks_repo_name: bubasooo/proggggggggg


text:
using Random
using LinearAlgebra
# const ContactImplicitMPC = Main

# include(joinpath(module_dir(), "src", "solver", "interior_point.jl"))
# include(joinpath(module_dir(), "src", "solver", "mehrotra.jl"))

################################################################################
# Test Utils
################################################################################

function get_initialization(ref_traj, t)
	Random.seed!(10)
	z = deepcopy(ref_traj.z[t])
	z += rand(length(z))
	θ = deepcopy(ref_traj.θ[t])
	return z, θ
end

function interior_point_timing(ref_traj, t, ip0)
	e1 = @belapsed begin
		z, θ = get_initialization(ref_traj, t)
		ip = deepcopy(ip0)
	end
	e2 = @belapsed begin
		z, θ = get_initialization(ref_traj, t)
		ip = deepcopy(ip0)
		interior_point_solve!(ip, z, θ)
	end
	@show e1
	@show e2
	return e2 - e1
end

# function interior_point_profiling(ref_traj, t, ip0, N::Int = 100)
# 	@profiler begin
# 		for i = 1:N
# 			z, θ = get_initialization(ref_traj, t)
# 			ip = deepcopy(ip0)
# 			interior_point_solve!(ip, z, θ)
# 		end
# 	end
# 	return nothing
# end

################################################################################
# Test data
################################################################################

s = ContactImplicitMPC.get_simulation("quadruped", "flat_2D_lc", "flat")
ref_traj = deepcopy(ContactImplicitMPC.get_trajectory(s.model, s.env,
    joinpath(ContactImplicitMPC.module_dir(), "src/dynamics/quadruped/gaits/gait2.jld2"),
    load_type = :split_traj_alt))

model = s.model
env = s.env
t = 10
nz = num_var(model, env)
nθ = num_data(model)

################################################################################
# Test Interior Point on the full non linear problem
################################################################################

z1, θ1 = get_initialization(ref_traj, t)
ip1 = interior_point(z1, θ1,
	oss = OptimizationSpace13(model, env),
	r! = s.res.r!,
    rz! = s.res.rz!,
    rθ! = s.res.rθ!,
    rz = s.rz,
    rθ = s.rθ,
    opts = InteriorPointOptions(
        max_iter=30,
        r_tol=1e-8,
        κ_tol=1e-8,
		# verbose=true
		))

interior_point_solve!(ip1)
r1 = zeros(nz)
s.res.r!(r1, ip1.z, ip1.θ, 0.0)
@testset "Interior Point Non Linear" begin
	@test norm(r1, Inf) < 1e-8
	@test abs(norm(r1, Inf) - 3.66e-9) < 1e-11
	@test ip1.iterations == 10
end
# e_ip = 1e6 * interior_point_timing(ref_traj, t, ip1)

################################################################################
# Test Interior Point on the linearized problem
################################################################################

im_traj1 = ImplicitTraj(ref_traj, s;
	κ = 1e-8,
	ip_type = :interior_point,
	opts = InteriorPointOptions(
			κ_tol = 1.0e-8,
			r_tol = 1.0e-8,
			diff_sol = true,
			max_iter=100,
			ϵ_min=0.05,
			solver=:empty_solver,
			# verbose=true
			))

z1, θ1 = get_initialization(ref_traj, t)
ip1 = deepcopy(im_traj1.ip[10])
interior_point_solve!(ip1, z1, θ1)
r1 = zeros(nz)
s.res.r!(r1, ip1.z, ip1.θ, 0.0)
@testset "Interior Point Linear" begin
	@test norm(r1, Inf) < 1e-8
end

# e_ip = 1e6 * interior_point_timing(ref_traj, t, im_traj1.ip[t])

# @profiler for k = 1:15000
# 	z1, θ1 = get_initialization(ref_traj, t)
# 	ip1 = deepcopy(im_traj1.ip[t])
# 	interior_point_solve!(ip1, z1, θ1)
# end

# z1, θ1 = get_initialization(ref_traj, t)
# ip1 = deepcopy(im_traj1.ip[t])
# @code_warntype interior_point_solve!(ip1)


text:
using ReachableControls
using Documenter

DocMeta.setdocmeta!(ReachableControls, :DocTestSetup, :(using ReachableControls); recursive=true)

makedocs(;
    modules=[ReachableControls],
    authors="Jonnie Diegelman <jonniediegelman@gmail.com> and contributors",
    repo="https://github.com/jonniedie/ReachableControls.jl/blob/{commit}{path}#{line}",
    sitename="ReachableControls.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://jonniedie.github.io/ReachableControls.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/jonniedie/ReachableControls.jl",
    devbranch="main",
)


text:
# Types for Boundary Conditions
    # Abstracts
    abstract type BoundaryCondition end
    abstract type HomogeneousBoundaryCondition <: BoundaryCondition end

# Types for Differences
    abstract type DifferenceMethod{N} end
    struct DefaultFirstDifference <: DifferenceMethod{1} end
    struct ForwardFirstDifference <: DifferenceMethod{1} end
    struct BackwardFirstDifference <: DifferenceMethod{1} end
    struct CentralSecondDifference <: DifferenceMethod{2} end

# Concretes
    struct Reflecting <: HomogeneousBoundaryCondition end
    @with_kw struct Mixed <: HomogeneousBoundaryCondition
        ξ::Real = 0.0
        direction = :auto # :forward/:backward/:auto
    end
    struct Absorbing <: HomogeneousBoundaryCondition end


text:
function g_pot(latitude, longitude, d_, t_)

standardmeridian=longitude

localstandardtime=t_
tthet=2*pi*(d_-1.0)/365.0

eqoftime=(0.000075+0.001868*cos(tthet)-0.032077*sin(tthet)-0.014615*cos(2*tthet)-0.040849*sin(2*tthet))*229.18

localapparentsolartime=localstandardtime+eqoftime/60+longitude/15
signedLAS=12-localapparentsolartime

signedLAS=abs(signedLAS)
omega= -15*signedLAS

decl_rad=0.006918-0.399912*cos(tthet)+0.070257*sin(tthet)-0.006758*cos(2*tthet)+0.000907*sin(2*tthet)-0.002697*cos(3*tthet)+0.00148*sin(3*tthet)
decl=decl_rad*180/pi
lat_rad=latitude*pi/180
long_rad=longitude*pi/180

theta_rad=acos(sin(decl_rad)*sin(lat_rad)+cos(decl_rad)*cos(lat_rad)*cos(omega*pi/180.))

theta=theta_rad/pi*180
beta=90-theta
beta_rad=beta*pi/180

dt=d_+t_/24

solarconst=1376

Rpot=solarconst*(1.00011+0.034221*cos(tthet)+0.00128*sin(tthet)+0.000719*cos(2*tthet)+0.000077*sin(2*tthet))

Rpot_h=Rpot*cos(theta_rad)
Rpot_h = cos(theta_rad)<0.0 ? 0.0 : Rpot_h
return(Rpot_h)
end


text:
#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/build_literate/Simulation_ContinuousWave.ipynb)

# # Continuous Wave Simulation
# The following code replicates the continuous wave simulation of Fig. 2 and is slightly more comprehensive in the sense that all discussed models are simulated.

# For these simulations we need the following packages:

using MRIgeneralizedBloch
using DifferentialEquations
using QuadGK
using Plots
plotlyjs(bg = RGBA(31/255,36/255,36/255,1.0), ticks=:native); #hide #!nb

# and we simulate an isolated semi-solid spin pool with the following parameters:
R₁ = 1.0 # 1/s
T₂ˢ = 10e-6 # s

version = "v2" #src
Tʳᶠ = 2e-3 # s
ω₁ = 2000π # rad/s
ω₀ = 200π # rad/s
#src version = "v1"
#src Tʳᶠ = 1 # s
#src ω₁ = 1e2 * 2π # rad/s

t = range(0, Tʳᶠ, length=1001) # time points for plotting
tspan = (0.0, Tʳᶠ); # simulation range

# These parameters correspond to Fig. 2b, the parameters for replicating Fig. 2a are `ω₁ = 200π` and `Tʳᶠ = 1`s.

# ## Lorentzian Lineshape
# In this script, we simulate the three lineshapes separately, starting with the Lorentzian lineshape for which the Bloch model provides a ground truth.

# ### Bloch Model
# We can formulate the [Bloch model](http://dx.doi.org/10.1103/PhysRev.70.460) as
# ```math
# \partial_t \begin{pmatrix} x \\ y \\ z \\ 1 \end{pmatrix} = \begin{pmatrix}
# -R_2 & -ω_0 & ω_1 & 0 \\
# ω_0 & -R_2 & 0 & 0 \\
# -ω_1 & 0 & -R_1 & R_1 \\
# 0 & 0 & 0 & 0
# \end{pmatrix} \begin{pmatrix} x \\ y \\ z \\ 1 \end{pmatrix} ,
# ```
# where the matrix is the Hamiltonian of the Bloch model. For a constant ``ω_0`` and ``ω_1``, we can evaluate the Bloch model by taking the  matrix exponential of its Hamiltonian:

H(ω₁, ω₀, R₂, R₁) = [-R₂  -ω₀  ω₁  0;
                       ω₀ -R₂   0  0;
                      -ω₁   0 -R₁ R₁;
                        0   0   0  0]

z_Bloch = similar(t)
for i = 1:length(t)
    (_, _, z_Bloch[i], _) = exp(H(ω₁, ω₀, 1 / T₂ˢ, R₁) * t[i]) * [0; 0; 1; 1]
end

# ### Henkelman's Steady-State Solution
# When assuming an isolated semi-solid pool, Eq. (9) in [*Henkelman, R. Mark, et al. "Quantitative interpretation of magnetization transfer." Magnetic resonance in medicine 29.6 (1993): 759-766*](https://doi.org/10.1002/mrm.1910290607) reduces to

g_Lorentzian(ω₀) = T₂ˢ / π / (1 + (T₂ˢ * ω₀)^2)
z_steady_state_Lorentzian = R₁ / (R₁ + π * ω₁^2 * g_Lorentzian(ω₀))

# where `g_Lorentzian(ω₀)` is the Lorentzian lineshape.

# ### Graham's Single Frequency Approximation
# The lineshape is also used to calculate [Graham's single frequency approximation](http://doi.org/10.1002/jmri.1880070520), which describes an exponential decay with the RF-induced saturation rate `Rʳᶠ`:

Rʳᶠ = π * ω₁^2 * g_Lorentzian(ω₀)
z_Graham_Lorentzian = @. (Rʳᶠ * exp(-t * (R₁ + Rʳᶠ)) + R₁) / (R₁ + Rʳᶠ);

# ### Sled's Model
# [Sled's model](http://dx.doi.org/10.1006/jmre.2000.2059) is given by the ordinary differential equation (ODE)
# ```math
# \partial_t z(t) = \left(-\pi \int_0^t G(t-τ) \omega_1(τ)^2 dτ \right) z(t)  + R_1 (1-z),
# ```
# where ``G(t-τ)`` is the Green's function. The Hamiltonian of this ODE is implemented in [`apply_hamiltonian_sled!`](@ref) and we solve this ODE with the [DifferentialEquations.jl](https://diffeq.sciml.ai/stable/) package:

z₀ = [1.0] # initial z-magnetization
param = (ω₁, 1, ω₀, R₁, T₂ˢ, greens_lorentzian) # defined by apply_hamiltonian_sled!
prob = ODEProblem(apply_hamiltonian_sled!, z₀, tspan, param)
z_Sled_Lorentzian = solve(prob);

# ### Generalized Bloch Model
# The generalized Bloch model is an integro-differential equation (IDE) as it depends on z(τ) instead of z(t):
# ```math
# \partial_t z(t) = - ω_1(t) \int_0^t G(t,τ) ω_1(τ) z(τ) dτ + R_1 (1 - z(t)) .
# ```
# For off-resonant RF-pulses with ``ω_1 = ω_x + i ω_y``, it is given by
# ```math
# \partial_t z(t) = - ω_y(t) \int_0^t G(t,τ) ω_y(τ) z(τ) dτ - ω_x(t) \int_0^t G(t,τ) ω_x(τ) z(τ) dτ + R_1 (1 - z(t)) .
# ```
# The Hamiltonian of the IDE is implemented in [`apply_hamiltonian_gbloch!`](@ref) and we can solve this IDE with the [delay-differential equation (DDE)](https://diffeq.sciml.ai/stable/tutorials/dde_example/) solver of the [DifferentialEquations.jl](https://diffeq.sciml.ai/stable/) package:

zfun(p, t) = [1.0] # initialize history function (will be populated with an interpolation by the DDE solver)

param = (ω₁, 1, ω₀, R₁, T₂ˢ, greens_lorentzian) # defined by apply_hamiltonian_gbloch!
prob = DDEProblem(apply_hamiltonian_gbloch!, z₀, zfun, tspan, param)
z_gBloch_Lorentzian = solve(prob);

# Now that we have solved all five models, we can plot the solutions for comparison:

p = plot(xlabel="t [ms]", ylabel="zˢ(t)")
plot!(p, 1e3t, zero(similar(t)) .+ z_steady_state_Lorentzian, label="Henkelman's steady-state")
plot!(p, 1e3t, z_Graham_Lorentzian, label="Graham's model")
plot!(p, 1e3t, (hcat(z_Sled_Lorentzian(t).u...)'), label="Sled's model")
plot!(p, 1e3t, (hcat(z_gBloch_Lorentzian(t).u...)'), label="generalized Bloch model")
plot!(p, 1e3t, z_Bloch, label="Bloch model")
#md Main.HTMLPlot(p) #hide

# Zooming into the plot, reveals virtually perfect (besides numerical differences) agreement between Bloch and generalized Bloch model and subtle, but existing differences when compared to the other models. Choosing a longer `T₂ˢ` amplifies these differences.

# ## Gaussian Lineshape
# We can repeat these simulations (with the exception of the Bloch model) for the Gaussian lineshape:

g_Gaussian(ω₀) = T₂ˢ / sqrt(2π) * exp(-(T₂ˢ * ω₀)^2 / 2)
z_steady_state_Gaussian = R₁ / (R₁ + π * ω₁^2 * g_Gaussian(ω₀))

Rʳᶠ = π * ω₁^2 * g_Gaussian(ω₀)
z_Graham_Gaussian = @. (Rʳᶠ * exp(-t * (R₁ + Rʳᶠ)) + R₁) / (R₁ + Rʳᶠ)

param = (ω₁, 1, ω₀, R₁, T₂ˢ, greens_gaussian) # defined by apply_hamiltonian_sled!
prob = ODEProblem(apply_hamiltonian_sled!, z₀, tspan, param)
z_Sled_Gaussian = solve(prob)

prob = DDEProblem(apply_hamiltonian_gbloch!, z₀, zfun, tspan, param)
z_gBloch_Gaussian = solve(prob)

p = plot(xlabel="t [ms]", ylabel="zˢ(t)")
plot!(p, 1e3t, zero(similar(t)) .+ z_steady_state_Gaussian, label="Henkelman's steady-state")
plot!(p, 1e3t, z_Graham_Gaussian, label="Graham' model")
plot!(p, 1e3t, (hcat(z_Sled_Gaussian(t).u...)'), label="Sled's model")
plot!(p, 1e3t, (hcat(z_gBloch_Gaussian(t).u...)'), label="generalized Bloch model")
#md Main.HTMLPlot(p) #hide

# ## Super-Lorentzian Lineshape
# And we can repeat these simulations (with the exception of the Bloch model) for the [super-Lorentzian lineshape](http://dx.doi.org/10.1002/mrm.1910330404), which reveals the most pronounced deviations between the models due to the substantially slower decay of the Green's function:

g_superLorentzian(ω₀) = sqrt(2 / π) * T₂ˢ * quadgk(ct -> exp(-2 * (T₂ˢ * ω₀ / abs(3 * ct^2 - 1))^2) / abs(3 * ct^2 - 1), 0.0, sqrt(1 / 3), 1)[1]
z_steady_state_superLorentzian = R₁ / (R₁ + π * ω₁^2 * g_superLorentzian(ω₀))

Rʳᶠ = π * ω₁^2 * g_superLorentzian(ω₀)
z_Graham_superLorentzian = @. (Rʳᶠ * exp(-t * (R₁ + Rʳᶠ)) + R₁) / (R₁ + Rʳᶠ)

G_superLorentzian = interpolate_greens_function(greens_superlorentzian, 0, Tʳᶠ/T₂ˢ)

param = (ω₁, 1, ω₀, R₁, T₂ˢ, G_superLorentzian)
prob = ODEProblem(apply_hamiltonian_sled!, z₀, tspan, param)
z_Sled_superLorentzian = solve(prob)

prob = DDEProblem(apply_hamiltonian_gbloch!, z₀, zfun, tspan, param)
z_gBloch_superLorentzian = solve(prob)

p = plot(xlabel="t [ms]", ylabel="zˢ(t)")
plot!(p, 1e3t, zero(similar(t)) .+ z_steady_state_superLorentzian, label="Henkelman's steady-state")
plot!(p, 1e3t, z_Graham_superLorentzian, label="Graham's model")
plot!(p, 1e3t, (hcat(z_Sled_superLorentzian(t).u...)'), label="Sled's model")
plot!(p, 1e3t, (hcat(z_gBloch_superLorentzian(t).u...)'), label="generalized Bloch model")
#md Main.HTMLPlot(p) #hide

#src ##################################################################
#src # export data for plotting
#src ##################################################################
using Printf #src
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/CW_Henkelman_steady_state_", version, ".txt")), "w") #src
write(io, "t_s Lorentzian Gaussian superLorentzian \n") #src
write(io, @sprintf("%1.3e %1.3e %1.3e %1.3e \n", 0, z_steady_state_Lorentzian, z_steady_state_Gaussian, z_steady_state_superLorentzian)) #src
write(io, @sprintf("%1.3e %1.3e %1.3e %1.3e \n", t[end], z_steady_state_Lorentzian, z_steady_state_Gaussian, z_steady_state_superLorentzian)) #src
close(io) #src

io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/CW_SpinDynamics_", version, ".txt")), "w") #src
write(io, "t_s t_ms z_Bloch z_gBloch_Lorentzian z_gBloch_Gaussian z_gBloch_superLorentzian z_Graham_Lorentzian z_Graham_Gaussian z_Graham_superLorentzian z_Sled_Lorentzian z_Sled_Gaussian z_Sled_superLorentzian \n") #src
for i = 1:30:length(t) #src
    write(io, @sprintf("%1.3e %1.3e %1.3e %1.3e %1.3e %1.3e %1.3e %1.3e %1.3e %1.3e %1.3e %1.3e \n", t[i], t[i]*1e3, z_Bloch[i], z_gBloch_Lorentzian(t[i])[1], z_gBloch_Gaussian(t[i])[1], z_gBloch_superLorentzian(t[i])[1], z_Graham_Lorentzian[i], z_Graham_Gaussian[i], z_Graham_superLorentzian[i], z_Sled_Lorentzian(t[i])[1], z_Sled_Gaussian(t[i])[1], z_Sled_superLorentzian(t[i])[1])) #src
end #src
close(io) #src


text:
using TimeModels

println("Testing Kalman.jl")

srand(1)

F = diagm([1.0])
V = diagm([2.0])
G = reshape([1, 2, -0.5], 3, 1)
W = diagm([8.0, 2.5, 4.0])
x0 = randn(1)
P0 = diagm([1e7])
mod1 = StateSpaceModel(F, V, G, W, x0, P0)

# Test simulating
println("Simulating series")
x, y = simulate(mod1, 100)
@time x, y = simulate(mod1, 100)
t = [1:size(x, 1)]

# Test filtering
println("Filtering")
filt = kalman_filter(y, mod1)
@time filt = kalman_filter(y, mod1)
println("Log likelihood:")
println(filt.loglik)

# Test smoothing
println("Smoothing")
smooth = kalman_smooth(y, mod1)
@time smooth = kalman_smooth(y, mod1)

# Test model-fitting
function build(theta)
	F = diagm(theta[1])
	V = diagm(exp(theta[2]))
	G = reshape(theta[3:5], 3, 1)
	W = diagm(exp(theta[6:8]))
	x0 = [theta[9]]
	P0 = diagm(1e7)
	StateSpaceModel(F, V, G, W, x0, P0)
end
theta0 = zeros(9)
build(theta0)

fit(y, build, theta0)

# println("Plotting")
# p = FramedPlot()
# for i=1:3
# 	add(p, Curve(t, y[:, i], "color", "grey"))
# end
# add(p, Curve(t, filt.filtered, "type", "dashed"))
# add(p, Curve(t, smooth.smoothed, "color", "blue"))
# # add(p, Curve(t, x[:, 1], "color", "blue"))
# # add(p, Curve(t, filt.filtered[:, 1], "color", "green"))
# # add(p, Curve(t, smooth.smoothed[:, 1], "color", "red"))
# Winston.display(p)

println("\n\nPassed Kalman tests.")


text:
function apply_to_neighbors(func, node; to_temp = false, only_to_outputs = false)
    if to_temp
        if only_to_outputs == false
            for in_node in node.temp_in_inhibitory
                func(in_node)
            end
            for in_node in node.temp_in_excitatory
                func(in_node)
            end
        end
        for out_node in node.temp_out_inhibitory
            func(out_node)
        end
        for out_node in node.temp_out_excitatory
            func(out_node)
        end
    else
        if only_to_outputs == false
            for in_node in node.in_inhibitory
                func(in_node)
            end
            for in_node in node.in_excitatory
                func(in_node)
            end
        end
        for out_node in node.out_inhibitory
            func(out_node)
        end
        for out_node in node.out_excitatory
            func(out_node)
        end
    end
end

function apply_to_all(func, network::Network; to_temp = false, only_to_outputs = false)
    already_visited_nodes = Set{Node}()

    for node in network.inputs
        _apply_to_all(func, node, already_visited_nodes, to_temp = to_temp, only_to_outputs = only_to_outputs)
    end
    for node in network.outputs
        _apply_to_all(func, node, already_visited_nodes, to_temp = to_temp, only_to_outputs = only_to_outputs)
    end
end
function apply_to_all(func, node::Node; to_temp = false, only_to_outputs = false)
    apply_to_all(func, [node], to_temp = to_temp, only_to_outputs = only_to_outputs)
end
function apply_to_all(func, nodes::Vector{Node}; to_temp = false, only_to_outputs = false)
    already_visited_nodes = Set{Node}()

    for node in nodes
        _apply_to_all(func, node, already_visited_nodes, to_temp = to_temp, only_to_outputs = only_to_outputs)
    end
end
function _apply_to_all(func, node, already_visited_nodes; to_temp = false, only_to_outputs = false)
    if node in already_visited_nodes
        return
    end

    push!(already_visited_nodes, node)
    func(node)

    apply_to_neighbors(node_ -> _apply_to_all(func, node_, already_visited_nodes, to_temp = to_temp, only_to_outputs = only_to_outputs), node,
                        to_temp = to_temp, only_to_outputs = only_to_outputs)
end

max_stars_repo_name: tochsner/ALifeBenchmark

max_issues_repo_name: tochsner/ALifeBenchmark

max_forks_repo_name: tochsner/ALifeBenchmark


text:
export Trajectory,
    SharedTrajectory,
    EpisodicTrajectory,
    CombinedTrajectory,
    CircularCompactSATrajectory,
    VectCompactSATrajectory,
    ElasticCompactSATrajectory,
    CircularCompactSALTrajectory,
    CircularCompactSARTSATrajectory,
    VectCompactSARTSATrajectory,
    ElasticCompactSARTSATrajectory,
    CircularCompactPSARTSATrajectory,
    CircularCompactSALRTSALTrajectory,
    CircularCompactPSALRTSALTrajectory

using MacroTools: @forward
using ElasticArrays

#####
# Trajectory
#####

"""
    Trajectory(;[trace_name=trace_container]...)

Simply a wrapper of `NamedTuple`.
Define our own type here to avoid type piracy with `NamedTuple`
"""
struct Trajectory{T} <: AbstractTrajectory
    traces::T
end

Trajectory(; kwargs...) = Trajectory(kwargs.data)

const DUMMY_TRAJECTORY = Trajectory()
const DummyTrajectory = typeof(DUMMY_TRAJECTORY)

@forward Trajectory.traces Base.keys, Base.haskey, Base.getindex

Base.push!(t::Trajectory, kv::Pair{Symbol}) = push!(t[first(kv)], last(kv))
Base.pop!(t::Trajectory, s::Symbol) = pop!(t[s])

isfull(t::Trajectory) = all(isfull, t.traces)

#####
# SharedTrajectory
#####

struct SharedTrajectoryMeta
    start_shift::Int
    end_shift::Int
end

"""
    SharedTrajectory(trace_container, meta::NamedTuple{([trace_name::Symbol],...), Tuple{[SharedTrajectoryMeta]...}})

Create multiple traces sharing the same underlying container.
"""
struct SharedTrajectory{X,M} <: AbstractTrajectory
    x::X
    meta::M
end

"""
    SharedTrajectory(trace_container, s::Symbol)

Automatically create the following three traces:

- `s`, share the data in `trace_container` in the range of `1:end-1`
- `s` with a prefix of `next_`, share the data in `trace_container` in the range of `2:end`
- `s` with a prefix of `full_`, a view of `trace_container`
"""
function SharedTrajectory(x, s::Symbol)
    SharedTrajectory(
        x,
        (;
            s => SharedTrajectoryMeta(1, -1),
            Symbol(:next_, s) => SharedTrajectoryMeta(2, 0),
            Symbol(:full_, s) => SharedTrajectoryMeta(1, 0),
        ),
    )
end

@forward SharedTrajectory.meta Base.keys, Base.haskey

function Base.getindex(t::SharedTrajectory, s::Symbol)
    m = t.meta[s]
    select_last_dim(t.x, m.start_shift:(nframes(t.x)+m.end_shift))
end

Base.push!(t::SharedTrajectory, kv::Pair{Symbol}) = push!(t.x, last(kv))
Base.empty!(t::SharedTrajectory) = empty!(t.x)
Base.pop!(t::SharedTrajectory, s::Symbol) = pop!(t.x)

function Base.pop!(t::SharedTrajectory)
    s = first(keys(t))
    (; s => pop!(t.x))
end

isfull(t::SharedTrajectory) = isfull(t.x)

#####
# EpisodicTrajectory
#####

"""
    EpisodicTrajectory(traces::T, flag_trace=:terminal)

Assuming that the `flag_trace` is in `traces` and it's an `AbstractVector{Bool}`,
meaning whether an environment reaches terminal or not. The last element in
`flag_trace` will be used to determine whether the whole trace is full or not.
"""
struct EpisodicTrajectory{T,flag_trace} <: AbstractTrajectory
    traces::T
end

EpisodicTrajectory(traces::T, flag_trace = :terminal) where {T} =
    EpisodicTrajectory{T,flag_trace}(traces)

@forward EpisodicTrajectory.traces Base.keys,
Base.haskey,
Base.getindex,
Base.push!,
Base.pop!,
Base.empty!

function isfull(t::EpisodicTrajectory{<:Any,F}) where {F}
    x = t.traces[F]
    (nframes(x) > 0) && select_last_frame(x)
end

#####
# CombinedTrajectory
#####

"""
    CombinedTrajectory(t1::AbstractTrajectory, t2::AbstractTrajectory)
"""
struct CombinedTrajectory{T1,T2} <: AbstractTrajectory
    t1::T1
    t2::T2
end

Base.haskey(t::CombinedTrajectory, s::Symbol) = haskey(t.t1, s) || haskey(t.t2, s)
Base.getindex(t::CombinedTrajectory, s::Symbol) =
    if haskey(t.t1, s)
        getindex(t.t1, s)
    elseif haskey(t.t2, s)
        getindex(t.t2, s)
    else
        throw(ArgumentError("unknown key: $s"))
    end

Base.keys(t::CombinedTrajectory) = (keys(t.t1)..., keys(t.t2)...)

Base.push!(t::CombinedTrajectory, kv::Pair{Symbol}) =
    if haskey(t.t1, first(kv))
        push!(t.t1, kv)
    elseif haskey(t.t2, first(kv))
        push!(t.t2, kv)
    else
        throw(ArgumentError("unknown kv: $kv"))
    end

Base.pop!(t::CombinedTrajectory, s::Symbol) =
    if haskey(t.t1, s)
        pop!(t.t1, s)
    elseif haskey(t.t2, s)
        pop!(t.t2, s)
    else
        throw(ArgumentError("unknown key: $s"))
    end

Base.pop!(t::CombinedTrajectory) = merge(pop!(t.t1), pop!(t.t2))

function Base.empty!(t::CombinedTrajectory)
    empty!(t.t1)
    empty!(t.t2)
end

isfull(t::CombinedTrajectory) = isfull(t.t1) && isfull(t.t2)

#####
# VectCompactSATrajectory
#####

const VectCompactSATrajectory = CombinedTrajectory{
    <:SharedTrajectory{<:Vector,<:NamedTuple{(:state, :next_state, :full_state)}},
    <:SharedTrajectory{<:Vector,<:NamedTuple{(:action, :next_action, :full_action)}},
}

function VectCompactSATrajectory(; state_type = Int, action_type = Int)
    CombinedTrajectory(
        SharedTrajectory(Vector{state_type}(), :state),
        SharedTrajectory(Vector{action_type}(), :action),
    )
end

#####
# CircularCompactSATrajectory
#####

const CircularCompactSATrajectory = CombinedTrajectory{
    <:SharedTrajectory{
        <:CircularArrayBuffer,
        <:NamedTuple{(:state, :next_state, :full_state)},
    },
    <:SharedTrajectory{
        <:CircularArrayBuffer,
        <:NamedTuple{(:action, :next_action, :full_action)},
    },
}

function CircularCompactSATrajectory(;
    capacity,
    state_type = Int,
    state_size = (),
    action_type = Int,
    action_size = (),
)
    CombinedTrajectory(
        SharedTrajectory(
            CircularArrayBuffer{state_type}(state_size..., capacity + 1),
            :state,
        ),
        SharedTrajectory(
            CircularArrayBuffer{action_type}(action_size..., capacity + 1),
            :action,
        ),
    )
end

#####
# ElasticCompactSATrajectory
#####

const ElasticCompactSATrajectory = CombinedTrajectory{
    <:SharedTrajectory{<:ElasticArray,<:NamedTuple{(:state, :next_state, :full_state)}},
    <:SharedTrajectory{<:ElasticArray,<:NamedTuple{(:action, :next_action, :full_action)}},
}

function ElasticCompactSATrajectory(;
    state_type = Int,
    state_size = (),
    action_type = Int,
    action_size = (),
)
    CombinedTrajectory(
        SharedTrajectory(ElasticArray{state_type}(undef, state_size..., 0), :state),
        SharedTrajectory(ElasticArray{action_type}(undef, action_size..., 0), :action),
    )
end

#####
# CircularCompactSALTrajectory
#####

const CircularCompactSALTrajectory = CombinedTrajectory{
    <:SharedTrajectory{
        <:CircularArrayBuffer,
        <:NamedTuple{
            (:legal_actions_mask, :next_legal_actions_mask, :full_legal_actions_mask),
        },
    },
    <:CircularCompactSATrajectory,
}

function CircularCompactSALTrajectory(;
    capacity,
    legal_actions_mask_size,
    legal_actions_mask_type = Bool,
    kw...,
)
    CombinedTrajectory(
        SharedTrajectory(
            CircularArrayBuffer{legal_actions_mask_type}(
                legal_actions_mask_size...,
                capacity + 1,
            ),
            :legal_actions_mask,
        ),
        CircularCompactSATrajectory(; capacity = capacity, kw...),
    )
end

#####
# VectCompactSARTSATrajectory
#####

const VectCompactSARTSATrajectory = CombinedTrajectory{
    <:Trajectory{<:NamedTuple{(:reward, :terminal),<:Tuple{<:Vector,<:Vector}}},
    <:VectCompactSATrajectory,
}

function VectCompactSARTSATrajectory(; reward_type = Float32, terminal_type = Bool, kw...)
    CombinedTrajectory(
        Trajectory(reward = Vector{reward_type}(), terminal = Vector{terminal_type}()),
        VectCompactSATrajectory(; kw...),
    )
end

#####
# CircularCompactSARTSATrajectory
#####

const CircularCompactSARTSATrajectory = CombinedTrajectory{
    <:Trajectory{
        <:NamedTuple{
            (:reward, :terminal),
            <:Tuple{<:CircularArrayBuffer,<:CircularArrayBuffer},
        },
    },
    <:CircularCompactSATrajectory,
}

function CircularCompactSARTSATrajectory(;
    capacity,
    reward_type = Float32,
    reward_size = (),
    terminal_type = Bool,
    terminal_size = (),
    kw...,
)
    CombinedTrajectory(
        Trajectory(
            reward = CircularArrayBuffer{reward_type}(reward_size..., capacity),
            terminal = CircularArrayBuffer{terminal_type}(terminal_size..., capacity),
        ),
        CircularCompactSATrajectory(; capacity = capacity, kw...),
    )
end

#####
# ElasticCompactSARTSATrajectory
#####

const ElasticCompactSARTSATrajectory = CombinedTrajectory{
    <:Trajectory{<:NamedTuple{(:reward, :terminal),<:Tuple{<:ElasticArray,<:ElasticArray}}},
    <:ElasticCompactSATrajectory,
}

function ElasticCompactSARTSATrajectory(;
    reward_type = Float32,
    reward_size = (),
    terminal_type = Bool,
    terminal_size = (),
    kw...,
)
    CombinedTrajectory(
        Trajectory(
            reward = ElasticArray{reward_type}(undef, reward_size..., 0),
            terminal = ElasticArray{terminal_type}(undef, terminal_size..., 0),
        ),
        ElasticCompactSATrajectory(; kw...),
    )
end

#####
# CircularCompactSALRTSALTrajectory
#####

const CircularCompactSALRTSALTrajectory = CombinedTrajectory{
    <:Trajectory{
        <:NamedTuple{
            (:reward, :terminal),
            <:Tuple{<:CircularArrayBuffer,<:CircularArrayBuffer},
        },
    },
    <:CircularCompactSALTrajectory,
}

function CircularCompactSALRTSALTrajectory(;
    capacity,
    reward_type = Float32,
    reward_size = (),
    terminal_type = Bool,
    terminal_size = (),
    kw...,
)
    CombinedTrajectory(
        Trajectory(
            reward = CircularArrayBuffer{reward_type}(reward_size..., capacity),
            terminal = CircularArrayBuffer{terminal_type}(terminal_size..., capacity),
        ),
        CircularCompactSALTrajectory(; capacity = capacity, kw...),
    )
end

#####
# CircularCompactPSARTSATrajectory
#####

const CircularCompactPSARTSATrajectory = CombinedTrajectory{
    <:Trajectory{
        <:NamedTuple{
            (:reward, :terminal, :priority),
            <:Tuple{<:CircularArrayBuffer,<:CircularArrayBuffer,<:SumTree},
        },
    },
    <:CircularCompactSATrajectory,
}

function CircularCompactPSARTSATrajectory(;
    capacity,
    priority_type = Float32,
    reward_type = Float32,
    reward_size = (),
    terminal_type = Bool,
    terminal_size = (),
    kw...,
)
    CombinedTrajectory(
        Trajectory(
            reward = CircularArrayBuffer{reward_type}(reward_size..., capacity),
            terminal = CircularArrayBuffer{terminal_type}(terminal_size..., capacity),
            priority = SumTree(priority_type, capacity),
        ),
        CircularCompactSATrajectory(; capacity = capacity, kw...),
    )
end

#####
# CircularCompactPSALRTSALTrajectory
#####

const CircularCompactPSALRTSALTrajectory = CombinedTrajectory{
    <:Trajectory{
        <:NamedTuple{
            (:reward, :terminal, :priority),
            <:Tuple{<:CircularArrayBuffer,<:CircularArrayBuffer,<:SumTree},
        },
    },
    <:CircularCompactSALTrajectory,
}

function CircularCompactPSALRTSALTrajectory(;
    capacity,
    priority_type = Float32,
    reward_type = Float32,
    reward_size = (),
    terminal_type = Bool,
    terminal_size = (),
    kw...,
)
    CombinedTrajectory(
        Trajectory(
            reward = CircularArrayBuffer{reward_type}(reward_size..., capacity),
            terminal = CircularArrayBuffer{terminal_type}(terminal_size..., capacity),
            priority = SumTree(priority_type, capacity),
        ),
        CircularCompactSALTrajectory(; capacity = capacity, kw...),
    )
end


text:
# Physical constants
const speed_of_light_mps = 299792458.0 # m/s
export speed_of_light_mps

# Limit data size/memory usage
const max_num_pixels_per_order = 8096
const max_num_orders = 128
const max_template_length = max_num_orders * max_num_pixels_per_order
const max_num_obs = 1024
const max_num_basis_vectors = 16
const min_num_star_basis_vectors = 1
const max_num_star_basis_vectors = max_num_basis_vectors
const min_num_telluric_basis_vectors = 0
const max_num_telluric_basis_vectors = max_num_basis_vectors
export max_num_pixels_per_order, max_num_orders, max_template_length, max_num_obs
export max_num_basis_vectors, min_num_star_basis_vectors, max_num_star_basis_vectors, min_num_telluric_basis_vectors, max_num_telluric_basis_vectors

# TODO: move to each instrument?
const default_blaze_degree = 4
const default_pixel_buffer_lo = 0
const default_pixel_buffer_hi = 0
export default_blaze_degree, default_pixel_buffer_lo, default_pixel_buffer_hi


text:
using VariationalInequality
using JuMP

using Test

@testset "Nagurney Model" begin
    include(joinpath(dirname(dirname(@__FILE__)), "example", "nagurney.jl"))
end

@testset "Traffic Equilibrium" begin
    include(joinpath(dirname(dirname(@__FILE__)), "example", "traffic.jl"))
end

@testset "Fukushima1" begin
    include(joinpath(dirname(dirname(@__FILE__)), "example", "fukushima1.jl"))
end


text:
using CommonLH
using Test

function data_settings_test()
    @testset "DataSettings" begin
		ds = CollegeStratData.test_data_settings();
		println(ds);
		println(CollegeStratData.settings_table(ds))
        @test n_school(ds) > 2
        @test n_colleges(ds) > 2
        @test n_2year(ds) < n_colleges(ds)
        @test n_gpa(ds) > 2
        @test n_parental(ds) > 2
        @test all(diff(hsgpa_ub(ds)) .> 0.0)
        @test all(diff(parental_ub(ds)) .> 0.0)
        @test all(hsgpa_masses(ds) .> 0.0)
        @test all(parental_masses(ds) .> 0.0)
        @test size(mass_gpa_yp(ds)) == (n_gpa(ds), n_parental(ds))

		@test is_two_year(ds, 1)
		@test !is_two_year(ds, 3)
		idx2V = two_year_colleges(ds);
		for ic in idx2V
			@test is_two_year(ds, ic)
		end

		idxNoV = no_grad_idx(ds);
		idxGradV = grad_idx(ds);
		@test vcat(idxNoV, idxGradV) == collect(1 : n_colleges(ds))

		g_edgeV = CollegeStratData.hsgpa_edges(ds);
		@test check_float_array(g_edgeV, 0.0, 1.0)
		g_massV = CollegeStratData.hsgpa_masses(ds)
		@test check_float_array(g_massV, 0.0, 0.5)
		@test isapprox(sum(g_massV), 1.0)
		@test isapprox(sum(parental_masses(ds)), 1.0)
		oneV = ones(length(g_massV))
		@test isapprox(mean_by_gpa(oneV, ds), 1.0)

		mass_gpM = mass_gpa_yp(ds);
		@test isapprox(sum(mass_gpM), 1.0)
		@test isapprox(vec(sum(mass_gpM, dims = 1)),  parental_masses(ds))
		@test isapprox(sum(mass_gpM, dims = 2),  hsgpa_masses(ds))

	end
end

@testset "DataSettings" begin
    data_settings_test()
end

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


text:
using JuMP
using DecisionProgramming
using RobustDecisionProgramming

function partial_expected_value(q, s_I, model, S, C, X, U⁺, z, k, hard_lower_bound, fixed)
    X′ = copy(X)
    X′[k].data[s_I..., :] = q
    P = DefaultPathProbability(C, X′)
    π_s = PathProbabilityVariables(model, z, S, P; hard_lower_bound=hard_lower_bound, fixed=fixed)
    pev = expected_value(model, π_s, U⁺)
    return pev
end

function minimum_partial_expected_value(dev, I_j, s_I, model, S, C, X, U⁺, z, k, hard_lower_bound)
    Q = polyhedral_uncertainty_set(dev)
    fixed = Dict{Node, State}(i=>s_i for (i, s_i) in zip(I_j, s_I))
    pevs = [partial_expected_value(q, s_I, model, S, C, X, U⁺, z, k, hard_lower_bound, fixed) for q in Q]
    x = @variable(model)
    for pev in pevs
        @constraint(model, x ≤ pev)
    end
    return (x=x, pevs=pevs, Q=Q, s_I=s_I)
end

"""Minimum expected value over discrete polyhedral uncertainty set."""
function min_expected_value(devs, model::Model, S::States, C::Vector{ChanceNode}, X, U⁺, z::DecisionVariables, k::Int; hard_lower_bound::Bool=false)
    mpevs = [minimum_partial_expected_value(dev, C[k].I_j, s_I, model, S, C, X, U⁺, z, k, hard_lower_bound) for (s_I, dev) in devs]
    min_ev = sum(mpev.x for mpev in mpevs)
    return (min_ev, mpevs)
end

"""Find distribution from uncertainty set that minimizes the partial expected value."""
function min_distribution(pevs, Q)
    return Q[argmin([value(pev) for pev in pevs])]
end

"""Consruct the probabilities that minimize the expected value."""
function min_probabilities(mpevs, X, k)
    X′ = copy(X)
    for (_, pevs, Q, s_I) in mpevs
        X′[k].data[s_I..., :] = min_distribution(pevs, Q)
    end
    return X′
end

max_stars_repo_name: jaantollander/RobustDecisionProgramming

max_issues_repo_name: jaantollander/RobustDecisionProgramming

max_forks_repo_name: jaantollander/RobustDecisionProgramming