text:
using Printf

include("workerPool.jl")

function cb(x::Int64)::Float64
        return 2.0
end

struct Box
	name::String
end

struct Results
	name::String
end

function cb_p2(x::Box)::Results
	Results(@sprintf("results-%s", Box.name))
end

p = WorkPool.Pool{Int64,Float64}(10,10,cb)

p2 = WorkPool.Pool{Box,Results}(10,10,cb_p2)

WorkPool.startPool(p)
WorkPool.push(p, WorkPool.WorkPkg{Int}("job1", 1))
sleep(2)
d = WorkPool.pop(p)
@info(@sprintf("job[%s] received %f\n", d.name, d.work))
WorkPool.shutdownPool(p)

WorkPool.startPool(p2)
WorkPool.push(p2, WorkPool.WorkPkg{Box}("foo", Box("box")))
sleep(2)
d = WorkPool.pop(p2)
@info(@sprintf("job[%s] received %s\n", d.name, d.work.name))
WorkPool.push(p2, WorkPool.WorkPkg{Box}("foo", Box("box")))
sleep(2)
d = WorkPool.pop(p2)
@info(@sprintf("job[%s] received %s\n", d.name, d.work.name))
WorkPool.push(p2, WorkPool.WorkPkg{Box}("foo", Box("box")))
sleep(2)
d = WorkPool.pop(p2)
@info(@sprintf("job[%s] received %s\n", d.name, d.work.name))
WorkPool.shutdownPool(p2)
WorkPool.getFinalState(p2)

max_stars_repo_name: rabarar/juliaSnippets

max_issues_repo_name: rabarar/juliaSnippets

max_forks_repo_name: rabarar/juliaSnippets


text:
using LinearAlgebra

struct LatticeVars
    # lattice vectors stored column-wise
    avec::Array{Float64,2}
    # magnitude of random displacements added to lattice vectors
    rndavec::Float64
    # inverse of lattice vector matrix
    ainv::Array{Float64,2}
    # reciprocal lattice vectors
    bvec::Array{Float64,2}
    # inverse of reciprocal lattice vector matrix
    binv::Array{Float64,2}
    # unit cell volume
    omega::Float64
    # Brillouin zone volume
    omegabz::Float64
    # any vector with length less than epslat is considered zero
    epslat::Float64
end

#=
case('avec')
  read(50,*,err=20) avec(:,1)
  read(50,*,err=20) avec(:,2)
  read(50,*,err=20) avec(:,3)

avec is stored like lattice vectors in Quantum Espresso
=#

function LatticeVars( LatVecs )
    @assert size(LatVecs,1) == 3
    @assert size(LatVecs,2) == 3

    avec = LatVecs[:,:]
    ainv = inv(avec)

    bvec = 2*pi*inv(avec')
    binv = inv(bvec)

    omega = det(LatVecs)
    omegabz = (2*pi)^3/omega

    epslat = 1e-6
    rndavec = 0.0

    return LatticeVars( avec, rndavec, ainv, bvec, binv, omega, omegabz, epslat )
end

ext: jl


text:
# This contains algorithms for breaking up a DAG into layers
"Calculate the layer of each node"
function layer_by_longest_path_to_source(graph, force_layer)
    dists = longest_paths(graph, sources(graph))
    force_layers!(graph, dists, force_layer)
    layer_groups = collect.(IterTools.groupby(i->dists[i], sort(vertices(graph), by=i->dists[i])))
    return layer_groups
end

"Correct the layer of each node, according to the optional parameter by user"
function force_layers!(graph, dists, force_layer::Vector{Pair{Int, Int}})
    # must process from end to beginning as otherwise can't move things after to make space
    ordered_forced_layers = sort(collect(force_layer), by=last; rev=true)
    for (node_id, target_layer) in ordered_forced_layers
        curr_layer = dists[node_id]
        if target_layer < curr_layer
            @warn "Ignored force_layer for node $node_id; curr layer ($curr_layer) > desired layer ($target_layer)"
        elseif any(dists[child] <= target_layer for child in outneighbors(graph, node_id))
            @warn "Ignored force_layer for node $node_id; as placing it at $target_layer would place it on same layer, or later than it's children."
        else
            dists[node_id] = target_layer
        end
    end
    return dists
end

###################
# helpers

function node2layer_lookup(layer2nodes)
    flat_map = [node=>layer_ind for (layer_ind, nodes) in enumerate(layer2nodes) for node in nodes]
    sort!(flat_map)
    @assert first.(flat_map) == 1:length(flat_map)
    return last.(sort(flat_map))
end

"add nodes so that no edges span multiple layers, returns mask"
function add_dummy_nodes!(graph, layer2nodes)
    dag_or_error(graph)
    nondummy_nodes = vertices(graph)
    # mapping from edges in the original graph to paths in the graph with dummy nodes
    edge_to_paths = Dict(e => eltype(graph)[] for e in edges(graph))
    node2layer = node2layer_lookup(layer2nodes)  # doesn't have dummy nodes, doesn't need them
    for cur_node in vertices(graph)
        cur_layer = node2layer[cur_node]
        for out_edge in filter(e -> src(e) == cur_node, collect(edges(graph)))  # need to copy as outwise will mutate when the graph is mutated
            out_node = out_edge.dst
            out_layer = node2layer[out_node]
            cur_layer < out_layer || throw(DomainError(node2layer, "Layer assigmenment must be strictly monotonic"))
            path = get!(edge_to_paths, out_edge, eltype(graph)[])
            if out_layer != cur_layer + 1
                rem_edge!(graph, cur_node, out_node) || error("removing edge failed")
                prev_node = cur_node
                push!(path, prev_node)
                for step_layer in (cur_layer + 1) : (out_layer - 1)
                    add_vertex!(graph) || throw(OverflowError("Could not add vertices to graph."))
                    step_node = nv(graph)
                    push!(layer2nodes[step_layer], step_node)
                    add_edge!(graph, prev_node, step_node)
                    push!(path, step_node)
                    prev_node = step_node
                end
                add_edge!(graph, prev_node, out_node)
                push!(path, out_node)
            else
                push!(path, cur_node, out_node)
            end
        end
    end
    is_dummy_mask = trues(nv(graph))
    is_dummy_mask[nondummy_nodes] .= false
    return is_dummy_mask, edge_to_paths
end


text:
using DiffEqSensitivity, Flux, Zygote, OrdinaryDiffEq, Test # , Plots

function lotka_volterra(du,u,p,t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = (α - β*y)x
  du[2] = dy = (δ*x - γ)y
end
p = [2.2, 1.0, 2.0, 0.4]
u0 = [1.0,1.0]
prob = ODEProblem(lotka_volterra,u0,(0.0,10.0),p)

# Reverse-mode

function predict_rd(p)
  Array(solve(prob,Tsit5(),p=p,saveat=0.1,reltol=1e-4,sensealg=TrackerAdjoint()))
end
loss_rd(p) = sum(abs2,x-1 for x in predict_rd(p))
loss_rd() = sum(abs2,x-1 for x in predict_rd(p))
loss_rd()

grads = Zygote.gradient(loss_rd, p)
@test !iszero(grads[1])

opt = ADAM(0.1)
cb = function ()
  display(loss_rd())
  # display(plot(solve(remake(prob,p=p),Tsit5(),saveat=0.1),ylim=(0,6)))
end

# Display the ODE with the current parameter values.
loss1 = loss_rd()
Flux.train!(loss_rd, Flux.params(p), Iterators.repeated((), 100), opt, cb = cb)
loss2 = loss_rd()
@test 10loss2 < loss1

# Forward-mode, R^n -> R^m layer

p = [2.2, 1.0, 2.0, 0.4]
function predict_fd()
  vec(Array(solve(prob,Tsit5(),p=p,saveat=0.0:0.1:1.0,reltol=1e-4,sensealg=ForwardDiffSensitivity())))
end
loss_fd() = sum(abs2,x-1 for x in predict_fd())
loss_fd()

ps = Flux.params(p)
grads = Zygote.gradient(loss_fd, ps)
@test !iszero(grads[p])

data = Iterators.repeated((), 100)
opt = ADAM(0.1)
cb = function ()
  display(loss_fd())
  # display(plot(solve(remake(prob,p=p),Tsit5(),saveat=0.1),ylim=(0,6)))
end

# Display the ODE with the current parameter values.
loss1 = loss_fd()
Flux.train!(loss_fd, ps, data, opt, cb = cb)
loss2 = loss_fd()
@test 10loss2 < loss1

# Adjoint sensitivity
p = [2.2, 1.0, 2.0, 0.4]
ps = Flux.params(p)
function predict_adjoint()
    solve(remake(prob,p=p),Tsit5(),saveat=0.1,reltol=1e-4)
end
loss_reduction(sol) = sum(abs2,x-1 for x in vec(sol))
loss_adjoint() = loss_reduction(predict_adjoint())
loss_adjoint()

grads = Zygote.gradient(loss_adjoint, ps)
@test !iszero(grads[p])

data = Iterators.repeated((), 100)
opt = ADAM(0.1)
cb = function ()
  display(loss_adjoint())
  # display(plot(solve(remake(prob,p=p),Tsit5(),saveat=0.1),ylim=(0,6)))
end

# Display the ODE with the current parameter values.
loss1 = loss_adjoint()
Flux.train!(loss_adjoint, ps, data, opt, cb = cb)
loss2 = loss_adjoint()
@test 10loss2 < loss1


text:
function agents_step!(agent, model)

end

function model_step!(model)
end


text:
# Dimensionality operations for 4D Tensors
# Author: Philipp Witte, pwitte3@gatech.edu
# Date: January 2020

export squeeze, unsqueeze, wavelet_squeeze, wavelet_unsqueeze, tensor_split, tensor_cat

####################################################################################################
# Squeeze and unsqueeze

"""
    Y = squeeze(X; pattern="column")

 Reshape input image such that each spatial dimension is reduced by a factor
 of 2, while the number of channels is increased by a factor of 4.

 *Input*:

 - `X`: 4D input tensor of dimensions `nx` x `ny` x `n_channel` x `batchsize`

 - `pattern`: Squeezing pattern

        1 2 3 4        1 1 3 3        1 3 1 3
        1 2 3 4        1 1 3 3        2 4 2 4
        1 2 3 4        2 2 4 4        1 3 1 3
        1 2 3 4        2 2 4 4        2 4 2 4

        column          patch       checkerboard

 *Output*:

 - `Y`: Reshaped tensor of dimensions `nx/2` x `ny/2` x `n_channel*4` x `batchsize`

 See also: [`unsqueeze`](@ref), [`wavelet_squeeze`](@ref), [`wavelet_unsqueeze`](@ref)
"""
function squeeze(X::AbstractArray{T,4}; pattern="column") where T

    # Dimensions
    nx_in, ny_in, nc_in, batchsize = size(X)
    if mod(nx_in, 2) == 1 || mod(ny_in, 2) == 1
        throw("Input dimensions must be multiple of 2")
    end
    nx_out = Int(round(nx_in/2))
    ny_out = Int(round(ny_in/2))
    nc_out = Int(round(nc_in*4))
    Y = cuzeros(X, nx_out, ny_out, nc_out, batchsize)

    if pattern == "column"
        Y = reshape(X, nx_out, ny_out, nc_out, batchsize)
    elseif pattern == "patch"
        Y[:, :, 1: nc_in, :] = X[1:nx_out, 1:ny_out, :, :]
        Y[:, :, nc_in+1: 2*nc_in, :] = X[nx_out+1:end, 1:ny_out, :, :]
        Y[:, :, 2*nc_in+1: 3*nc_in, :] = X[1:nx_out, ny_out+1:end, :, :]
        Y[:, :, 3*nc_in+1: 4*nc_in, :] = X[nx_out+1:end, ny_out+1:end, :, :]
    elseif pattern == "checkerboard"
        Y[:, :, 1: nc_in, :] = X[1:2:nx_in, 1:2:ny_in, 1:nc_in, :]
        Y[:, :, nc_in+1: 2*nc_in, :] = X[2:2:nx_in, 1:2:ny_in, 1:nc_in, :]
        Y[:, :, 2*nc_in+1: 3*nc_in, :] = X[1:2:nx_in, 2:2:ny_in, 1:nc_in, :]
        Y[:, :, 3*nc_in+1: 4*nc_in, :] = X[2:2:nx_in, 2:2:ny_in, 1:nc_in, :]
    else
        throw("Specified pattern not defined.")
    end
    return Y
end

function squeeze(X::AbstractArray{T,5}; pattern="column") where T

    # Dimensions
    nx_in, ny_in, nz_in, nc_in, batchsize = size(X)
    if mod(nx_in, 2) == 1 || mod(ny_in, 2) == 1
        throw("Input dimensions must be multiple of 2")
    end
    nx_out = Int(round(nx_in/2))
    ny_out = Int(round(ny_in/2))
    nz_out = Int(round(nz_in/2))
    nc_out = Int(round(nc_in*8))
    Y = cuzeros(X, nx_out, ny_out, nz_out, nc_out, batchsize)

    if pattern == "column"
        Y = reshape(X, nx_out, ny_out, nz_out, nc_out, batchsize)
    elseif pattern == "patch"
        Y[:, :, :, 1: nc_in, :] = X[1:nx_out, 1:ny_out, 1:nz_out, :, :]
        Y[:, :, :, nc_in+1: 2*nc_in, :] = X[nx_out+1:end, 1:ny_out, 1:nz_out, :, :]
        Y[:, :, :, 2*nc_in+1: 3*nc_in, :] = X[1:nx_out, ny_out+1:end, 1:nz_out, :, :]
        Y[:, :, :, 3*nc_in+1: 4*nc_in, :] = X[nx_out+1:end, ny_out+1:end, 1:nz_out, :, :]
        Y[:, :, :, 4*nc_in+1: 5*nc_in, :] = X[1:nx_out, 1:ny_out, nz_out+1:end, :, :]
        Y[:, :, :, 5*nc_in+1: 6*nc_in, :] = X[nx_out+1:end, 1:ny_out, nz_out+1:end, :, :]
        Y[:, :, :, 6*nc_in+1: 7*nc_in, :] = X[1:nx_out, ny_out+1:end, nz_out+1:end, :, :]
        Y[:, :, :, 7*nc_in+1: 8*nc_in, :] = X[nx_out+1:end, ny_out+1:end, nz_out+1:end, :, :]
    elseif pattern == "checkerboard"
        throw("Currently not implemented for 5D tensors.")
    else
        throw("Specified pattern not defined.")
    end
    return Y
end

"""
    X = unsqueeze(Y; pattern="column")

 Undo squeezing operation by reshaping input image such that each spatial dimension is
 increased by a factor of 2, while the number of channels is decreased by a factor of 4.

 *Input*:

 - `Y`: 4D input tensor of dimensions `nx` x `ny` x `n_channel` x `batchsize`

 - `pattern`: Squeezing pattern

            1 2 3 4        1 1 3 3        1 3 1 3
            1 2 3 4        1 1 3 3        2 4 2 4
            1 2 3 4        2 2 4 4        1 3 1 3
            1 2 3 4        2 2 4 4        2 4 2 4

            column          patch       checkerboard

 *Output*:

 - `X`: Reshaped tensor of dimensions `nx*2` x `ny*2` x `n_channel/4` x `batchsize`

 See also: [`squeeze`](@ref), [`wavelet_squeeze`](@ref), [`wavelet_unsqueeze`](@ref)
"""
function unsqueeze(Y::AbstractArray{T,4}; pattern="column") where T

    # Dimensions
    nx_in, ny_in, nc_in, batchsize = size(Y)
    if mod(nx_in, 2) == 1 || mod(ny_in, 2) == 1
        throw("Input dimensions must be multiple of 2")
    end
    nx_out = Int(round(nx_in*2))
    ny_out = Int(round(ny_in*2))
    nc_out = Int(round(nc_in/4))
    X = cuzeros(Y, nx_out, ny_out, nc_out, batchsize)

    if pattern == "column"
        X = reshape(Y, nx_out, ny_out, nc_out, batchsize)
    elseif pattern == "patch"
        X[1:nx_in, 1:ny_in, :, :] = Y[:, :, 1: nc_out, :]
        X[nx_in+1:end, 1:ny_in, :, :] = Y[:, :, nc_out+1: 2*nc_out, :]
        X[1:nx_in, ny_in+1:end, :, :] = Y[:, :, 2*nc_out+1: 3*nc_out, :]
        X[nx_in+1:end, ny_in+1:end, :, :] = Y[:, :, 3*nc_out+1: 4*nc_out, :]
    elseif pattern == "checkerboard"
        X[1:2:nx_out, 1:2:ny_out, 1:nc_out, :] = Y[:, :, 1: nc_out, :]
        X[2:2:nx_out, 1:2:ny_out, 1:nc_out, :] = Y[:, :, nc_out+1: 2*nc_out, :]
        X[1:2:nx_out, 2:2:ny_out, 1:nc_out, :] = Y[:, :, 2*nc_out+1: 3*nc_out, :]
        X[2:2:nx_out, 2:2:ny_out, 1:nc_out, :] = Y[:, :, 3*nc_out+1: 4*nc_out, :]
    else
        throw("Specified pattern not defined.")
    end
    return X
end

function unsqueeze(Y::AbstractArray{T,5}; pattern="column") where T

    # Dimensions
    nx_in, ny_in, nz_in, nc_in, batchsize = size(Y)
    if mod(nx_in, 2) == 1 || mod(ny_in, 2) == 1
        throw("Input dimensions must be multiple of 2")
    end
    nx_out = Int(round(nx_in*2))
    ny_out = Int(round(ny_in*2))
    nz_out = Int(round(ny_in*2))
    nc_out = Int(round(nc_in/8))
    X = cuzeros(Y, nx_out, ny_out, nz_out, nc_out, batchsize)

    if pattern == "column"
        X = reshape(Y, nx_out, ny_out, nz_out, nc_out, batchsize)
    elseif pattern == "patch"
        X[1:nx_in, 1:ny_in, 1:nz_in, :, :] = Y[:, :, :, 1: nc_out, :]
        X[nx_in+1:end, 1:ny_in, 1:nz_in, :, :] = Y[:, :, :, nc_out+1: 2*nc_out, :]
        X[1:nx_in, ny_in+1:end, 1:nz_in, :, :] = Y[:, :, :, 2*nc_out+1: 3*nc_out, :]
        X[nx_in+1:end, ny_in+1:end, 1:nz_in, :, :] = Y[:, :, :, 3*nc_out+1: 4*nc_out, :]
        X[1:nx_in, 1:ny_in, nz_in+1:end, :, :] = Y[:, :, :, 4*nc_out+1: 5*nc_out, :]
        X[nx_in+1:end, 1:ny_in, nz_in+1:end, :, :] = Y[:, :, :, 5*nc_out+1: 6*nc_out, :]
        X[1:nx_in, ny_in+1:end, nz_in+1:end, :, :] = Y[:, :, :, 6*nc_out+1: 7*nc_out, :]
        X[nx_in+1:end, ny_in+1:end, nz_in+1:end, :, :] = Y[:, :, :, 7*nc_out+1: 8*nc_out, :]
    elseif pattern == "checkerboard"
        throw("Specified pattern not defined.")
    else
        throw("Specified pattern not defined.")
    end
    return X
end

function unsqueeze(X::AbstractArray{T,4}, Y::AbstractArray{T,4}; pattern="column") where T
    return unsqueeze(X; pattern=pattern), unsqueeze(Y; pattern=pattern)
end

function unsqueeze(X::AbstractArray{T,5}, Y::AbstractArray{T,5}; pattern="column") where T
    return unsqueeze(X; pattern=pattern), unsqueeze(Y; pattern=pattern)
end

####################################################################################################
# Squeeze and unsqueeze using the wavelet transform

"""
    Y = wavelet_squeeze(X; type=WT.db1)

 Perform a 1-level channelwise 2D wavelet transform of X and squeeze output of each
 transform into 4 channels (per 1 input channel).

 *Input*:

 - `X`: 4D input tensor of dimensions `nx` x `ny` x `n_channel` x `batchsize`

 - `type`: Wavelet filter type. Possible values are `WT.haar` for Haar wavelets,
    `WT.coif2`, `WT.coif4`, etc. for Coiflet wavelets, or `WT.db1`, `WT.db2`, etc.
    for Daubechies wavetlets. See *https://github.com/JuliaDSP/Wavelets.jl* for a
    full list.

 *Output*:

 - `Y`: Reshaped tensor of dimensions `nx/2` x `ny/2` x `n_channel*4` x `batchsize`

 See also: [`wavelet_unsqueeze`](@ref), [`squeeze`](@ref), [`unsqueeze`](@ref)
"""
function wavelet_squeeze(X::AbstractArray{T,4}; type=WT.db1) where T
    nx_in, ny_in, nc_in, batchsize = size(X)
    nx_out = Int(round(nx_in/2))
    ny_out = Int(round(ny_in/2))
    nc_out = Int(round(nc_in*4))
    Y = cuzeros(X, nx_out, ny_out, nc_out, batchsize)
    for i=1:batchsize
        for j=1:nc_in
            Ycurr = dwt(X[:,:,j,i], wavelet(type), 1)
            Y[:, :, (j-1)*4 + 1: j*4, i] = squeeze(reshape(Ycurr, nx_in, ny_in, 1, 1); pattern="patch")
        end
    end
    return Y
end

function wavelet_squeeze(X::AbstractArray{T,5}; type=WT.db1) where T
    nx_in, ny_in, nz_in, nc_in, batchsize = size(X)
    nx_out = Int(round(nx_in/2))
    ny_out = Int(round(ny_in/2))
    nz_out = Int(round(nz_in/2))
    nc_out = Int(round(nc_in*8))
    Y = cuzeros(X, nx_out, ny_out, nz_out, nc_out, batchsize)
    for i=1:batchsize
        for j=1:nc_in
            Ycurr = dwt(X[:,:,:,j,i], wavelet(type), 1)
            Y[:, :, :, (j-1)*8 + 1: j*8, i] = squeeze(reshape(Ycurr, nx_in, ny_in, nz_in, 1, 1); pattern="patch")
        end
    end
    return Y
end

"""
    X = wavelet_unsqueeze(Y; type=WT.db1)

 Perform a 1-level inverse 2D wavelet transform of Y and unsqueeze output.
 This reduces the number of channels by 4 and increases each spatial
 dimension by a factor of 2. Inverse operation of `wavelet_squeeze`.

 *Input*:

 - `Y`: 4D input tensor of dimensions `nx` x `ny` x `n_channel` x `batchsize`

 - `type`: Wavelet filter type. Possible values are `haar` for Haar wavelets,
  `coif2`, `coif4`, etc. for Coiflet wavelets, or `db1`, `db2`, etc. for Daubechies
  wavetlets. See *https://github.com/JuliaDSP/Wavelets.jl* for a full list.

 *Output*:

 - `X`: Reshaped tensor of dimenions `nx*2` x `ny*2` x `n_channel/4` x `batchsize`

 See also: [`wavelet_squeeze`](@ref), [`squeeze`](@ref), [`unsqueeze`](@ref)
"""
function wavelet_unsqueeze(Y::AbstractArray{T,4}; type=WT.db1) where T
    nx_in, ny_in, nc_in, batchsize = size(Y)
    nx_out = Int(round(nx_in*2))
    ny_out = Int(round(ny_in*2))
    nc_out = Int(round(nc_in/4))
    X = cuzeros(Y, nx_out, ny_out, nc_out, batchsize)
    for i=1:batchsize
        for j=1:nc_out
            Ycurr = unsqueeze(Y[:, :, (j-1)*4 + 1: j*4, i:i]; pattern="patch")[:, :, 1, 1]
            Xcurr = idwt(Ycurr, wavelet(type), 1)
            X[:, :, j, i] = reshape(Xcurr, nx_out, ny_out, 1, 1)
        end
    end
    return X
end

function wavelet_unsqueeze(Y::AbstractArray{T,5}; type=WT.db1) where T
    nx_in, ny_in, nz_in, nc_in, batchsize = size(Y)
    nx_out = Int(round(nx_in*2))
    ny_out = Int(round(ny_in*2))
    nz_out = Int(round(nz_in*2))
    nc_out = Int(round(nc_in/8))
    X = cuzeros(Y, nx_out, ny_out, nz_out, nc_out, batchsize)
    for i=1:batchsize
        for j=1:nc_out
            Ycurr = unsqueeze(Y[:, :, :, (j-1)*8 + 1: j*8, i:i]; pattern="patch")[:, :, :, 1, 1]
            Xcurr = idwt(Ycurr, wavelet(type), 1)
            X[:, :, :, j, i] = reshape(Xcurr, nx_out, ny_out, nz_out, 1, 1)
        end
    end
    return X
end

####################################################################################################
# Split and concatenate

"""
    Y, Z = tensor_split(X)

 Split 4D input tensor in half along the channel dimension. Inverse operation
 of `tensor_cat`.

 *Input*:

 - `X`: 4D input tensor of dimensions `nx` x `ny` x `n_channel` x `batchsize`

 *Output*:

 - `Y`, `Z`: 4D output tensors, each of dimensions `nx` x `ny` x `n_channel/2` x `batchsize`

 See also: [`tensor_cat`](@ref)
"""
function tensor_split(X::AbstractArray{T,4}; split_index=nothing) where T
    if isnothing(split_index)
        k = Int(round(size(X, 3)/2))
    else
        k = split_index
    end
    return X[:, :, 1:k, :], X[:, :, k+1:end, :]
end

function tensor_split(X::AbstractArray{T,5}; split_index=nothing) where T
    if isnothing(split_index)
        k = Int(round(size(X, 4)/2))
    else
        k = split_index
    end
    return X[:, :, :, 1:k, :], X[:, :, :, k+1:end, :]
end

function tensor_split(X::AbstractArray{T,1}; split_index=nothing) where T
    if isnothing(split_index)
        k = Int(round(size(X, 1)/2))
    else
        k = split_index
    end
    return X[1:k], X[k+1:end]
end

"""
    X = tensor_cat(Y, Z)

 Concatenate 4D input tensors along the channel dimension. Inverse operation
 of `tensor_split`.

 *Input*:

 - `Y`, `Z`: 4D input tensors, each of dimensions `nx` x `ny` x `n_channel` x `batchsize`

 *Output*:

 - `X`: 4D output tensor of dimensions `nx` x `ny` x `n_channel*2` x `batchsize`

 See also: [`tensor_split`](@ref)
"""
function tensor_cat(X::AbstractArray{T,4}, Y::AbstractArray{T,4}) where T
    if size(X, 3) == 0
        return Y
    elseif size(Y, 3) == 0
        return X
    else
        return cat(X, Y; dims=3)
    end
end

function tensor_cat(X::AbstractArray{T,5}, Y::AbstractArray{T,5}) where T
    if size(X, 4) == 0
        return Y
    elseif size(Y, 4) == 0
        return X
    else
        return cat(X, Y; dims=4)
    end
end

function tensor_cat(X::AbstractArray{T,1}, Y::AbstractArray{T,1}) where T
    if size(X, 1) == 0
        return Y
    elseif size(Y, 1) == 0
        return X
    else
        return cat(X, Y; dims=1)
    end
end

tensor_cat(X::Tuple{AbstractArray{T,4}, AbstractArray{T,4}}) where T = tensor_cat(X[1], X[2])
tensor_cat(X::Tuple{AbstractArray{T,5}, AbstractArray{T,5}}) where T = tensor_cat(X[1], X[2])


text:
using ArgParse
using YAML
using Distributed

# const IN_SLURM = "SLURM_JOBID" in keys(ENV)
# IN_SLURM && using ClusterManagers
# if IN_SLURM
#     pids = addprocs(SlurmManager(parse(Int, ENV["SLURM_NTASKS"])))
# else
#     pids = addprocs()
# end

@everywhere include("evo/data.jl")

s = ArgParse.ArgParseSettings()

ArgParse.@add_arg_table(
    s,
    "--seed", arg_type=Int, default=0,
    "--id", arg_type=String, default="test",
    "--log", arg_type=String, default="evolution.log",
    "--cfg", arg_type=String, default="cfg/darwin.yaml",
)

args = ArgParse.parse_args(s)
cfg = YAML.load_file(args["cfg"])
cfg["seed"] = args["seed"]
cfg["n_fitness"] = 1

e = Evolution(NeurodevoInd, cfg; id=args["id"], logfile=args["log"])
e.mutation = uniform_mutation
e.evaluation = data_evaluation
e.generation = generation
run!(e)


text:
module Octrees

# Imports
import Base: *, display, getindex, iterate, in, lastindex
using RegionTrees
import RegionTrees: AbstractRefinery, needs_refinement, refine_data
using StaticArrays
using LinearAlgebra
using Random

include("types.jl")
include("codes.jl")
include("octree.jl")

export code_to_cell
export knn
export LocationalCode
export Octree

end # module


text:
# Autogenerated wrapper script for Cbc_jll for powerpc64le-linux-gnu-libgfortran5-cxx11
export cbc, libCbc, libOsiCbc, libcbcsolver

using ASL_jll
using Cgl_jll
using Clp_jll
using Osi_jll
using CoinUtils_jll
using OpenBLAS32_jll
using CompilerSupportLibraries_jll
JLLWrappers.@generate_wrapper_header("Cbc")
JLLWrappers.@declare_library_product(libCbc, "libCbc.so.3")
JLLWrappers.@declare_library_product(libOsiCbc, "libOsiCbc.so.3")
JLLWrappers.@declare_library_product(libcbcsolver, "libCbcSolver.so.3")
JLLWrappers.@declare_executable_product(cbc)
function __init__()
    JLLWrappers.@generate_init_header(ASL_jll, Cgl_jll, Clp_jll, Osi_jll, CoinUtils_jll, OpenBLAS32_jll, CompilerSupportLibraries_jll)
    JLLWrappers.@init_library_product(
        libCbc,
        "lib/libCbc.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

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

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

    JLLWrappers.@init_executable_product(
        cbc,
        "bin/cbc",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
# Composition
# ===========
#
# Sequence composition counter.
#
# This file is a part of BioJulia.
# License is MIT: https://github.com/BioJulia/Bio.jl/blob/master/LICENSE.md

type Composition{T} <: Associative{T,Int}
    counts::Vector{Int}
end

function Composition{A<:Union{DNAAlphabet,RNAAlphabet}}(seq::BioSequence{A})
    counts = zeros(Int, 16)
    @inbounds for x in seq
        counts[reinterpret(UInt8, x) + 1] += 1
    end
    return Composition{eltype(A)}(counts)
end

function Composition(seq::ReferenceSequence)
    counts = zeros(Int, 16)
    @inbounds for x in seq
        counts[reinterpret(UInt8, x) + 1] += 1
    end
    return Composition{DNANucleotide}(counts)
end

function Composition{T,k}(kmer::Kmer{T,k})
    counts = zeros(Int, 16)
    @inbounds begin
        counts[Int(DNA_A)+1] = count_a(kmer)
        counts[Int(DNA_C)+1] = count_c(kmer)
        counts[Int(DNA_G)+1] = count_g(kmer)
        counts[Int(DNA_T)+1] = count_t(kmer)  # U when T == RNANucleotide
    end
    return Composition{T}(counts)
end

function Composition(seq::AminoAcidSequence)
    counts = zeros(Int, length(alphabet(AminoAcid)))
    @inbounds for x in seq
        counts[reinterpret(UInt8, x) + 1] += 1
    end
    return Composition{AminoAcid}(counts)
end

function Composition{T,k}(iter::EachKmerIterator{T,k})
    counts = zeros(Int, 4^k)
    for (_, x) in iter
        counts[convert(UInt64, x) + 1] += 1
    end
    return Composition{Kmer{T,k}}(counts)
end

"""
    composition(seq | kmer_iter)

Calculate composition of biological symbols in `seq` or k-mers in `kmer_iter`.
"""
function composition(iter::Union{Sequence,EachKmerIterator})
    return Composition(iter)
end

function Base.:(==){T}(x::Composition{T}, y::Composition{T})
    return x.counts == y.counts
end

function Base.length(comp::Composition)
    return length(comp.counts)
end

function Base.start(comp::Composition)
    return UInt64(0)
end

function Base.done(comp::Composition, i)
    return i ≥ length(comp.counts)
end

function Base.next(comp::Composition, i)
    key = convert(keytype(comp), i)
    count = comp.counts[i + 1]
    return (key => count), i + 1
end

function Base.getindex{T}(comp::Composition{T}, x)
    i = convert(UInt64, convert(T, x))
    if !(0 ≤ i < endof(comp.counts))
        throw(KeyError(x))
    end
    return comp.counts[i+1]
end

function Base.copy{T}(comp::Composition{T})
    return Composition{T}(copy(comp.counts))
end

function Base.merge{T}(comp::Composition{T}, other::Composition{T})
    return merge!(copy(comp), other)
end

function Base.merge!{T}(comp::Composition{T}, other::Composition{T})
    @assert length(comp.counts) == length(other.counts)
    for i in 1:endof(comp.counts)
        comp.counts[i] += other.counts[i]
    end
    return comp
end

function Base.summary{T}(::Composition{T})
    if T == DNANucleotide
        return "DNA Nucleotide Composition"
    elseif T == RNANucleotide
        return "RNA Nucleotide Composition"
    elseif T == AminoAcid
        return "Amino Acid Composition"
    else
        return string(T, " Composition")
    end
end


text:
Base.show{T <: NumberParameter}(io::IO, n::T) = begin
    @printf io "    ["
    print_with_color(:blue, io, "NumberParameter")
    @printf io "]\n"

    print_with_color(:blue, io, "    name")
    @printf io " : "
    print_with_color(:bold, io, "$(n.name)\n")

    print_with_color(:green, io, "    min")
    @printf io "  : %6f\n" n.min

    print_with_color(:red, io, "    max")
    @printf io "  : %6f\n" n.max

    print_with_color(:yellow, io, "    value")
    @printf io ": %6f\n" n.value
    return
end

Base.show{T <: EnumParameter}(io::IO, n::T) = begin
    @printf io "    ["
    print_with_color(:blue, io, "EnumParameter")
    @printf io "]\n"
    print_with_color(:blue, io, "    name")
    @printf io "   : "
    print_with_color(:bold, io, "$(n.name)\n")
    print_with_color(:yellow, io, "    current")
    @printf io ": "
    print_with_color(:bold, io, "$(n.current.name)\n")
    print_with_color(:green, io, "    values")
    @printf io " : \n"
    for i = 1:length(n.values)
        show(io, n.values[i])
        if i < length(n.values) && length(n.values) > 1
            print_with_color(:blue, io, "    ***\n")
        end
    end
    return
end

Base.show{T <: PermutationParameter}(io::IO, n::T) = begin
    @printf io "    ["
    print_with_color(:blue, io, "PermutationParameter")
    @printf io "]\n"

    print_with_color(:blue, io, "    name")
    @printf io " : "
    print_with_color(:bold, io, "$(n.name)\n")

    print_with_color(:yellow, io, "    value")
    @printf io ": "
    show(io, n.value)
    @printf io "\n"

    print_with_color(:green, io, "    size")
    @printf io " : %d\n" n.size
    return
end

Base.show{T <: StringParameter}(io::IO, n::T) = begin
    @printf io "    ["
    print_with_color(:blue, io, "StringParameter")
    @printf io "]\n"
    print_with_color(:blue, io, "    name")
    @printf io " : "
    print_with_color(:bold, io, "$(n.name)\n")
    print_with_color(:yellow, io, "    value")
    @printf io ": \"%s\"\n" n.value
    return
end

Base.show{T <: BoolParameter}(io::IO, n::T) = begin
    @printf io "    ["
    print_with_color(:blue, io, "BoolParameter")
    @printf io "]\n"
    print_with_color(:blue, io, "    name")
    @printf io " : "
    print_with_color(:bold, io, "$(n.name)\n")
    print_with_color(:yellow, io, "    value")
    @printf io ": %s\n" n.value
    return
end

Base.show{T <: Configuration}(io::IO, n::T) = begin
    @printf io "  ["
    print_with_color(:blue, io, "Configuration")
    @printf io "]\n"
    print_with_color(:blue, io, "  name")
    @printf io "      : "
    print_with_color(:bold, io, "$(n.name)\n")
    print_with_color(:blue, io, "  parameters")
    @printf io ":\n"
    p = collect(keys(n.parameters))
    for i = 1:length(p)
        show(io, n[p[i]])
        if i < length(p) && length(p) > 1
            print_with_color(:blue, io, "    ***\n")
        end
    end
    return
end

Base.show{T <: Result}(io::IO, n::T) = begin
    if n.is_final
        @printf io "["
        print_with_color(:blue, io, "Final Result")
        @printf io "]\n"
        print_with_color(:yellow, io, "Cost")
        @printf io "                  : "
        print_with_color(:bold, io, "$(n.cost_minimum)\n")
        print_with_color(:yellow, io, "Found in Iteration")
        @printf io "    : "
        print_with_color(:bold, io, "$(n.iterations)\n")
        print_with_color(:blue, io, "Current Iteration")
        @printf io "     : "
        print_with_color(:bold, io, "$(n.current_iteration)\n")
        print_with_color(:blue, io, "Technique")
        @printf io "             : "
        print_with_color(:bold, io, "$(n.technique)\n")
        print_with_color(:blue, io, "Function Calls")
        @printf io "        : "
        print_with_color(:bold, io, "$(n.cost_calls)\n")
        print_with_color(:blue, io, "Starting Configuration")
        @printf io ":\n"
        show(io, n.start)
        print_with_color(:blue, io, "Minimum Configuration")
        @printf io " :\n"
        show(io, n.minimum)
    else
        @printf io "["
        print_with_color(:blue, io, "Result")
        @printf io "]\n"
        print_with_color(:yellow, io, "Cost")
        @printf io "              : "
        print_with_color(:bold, io, "$(n.cost_minimum)\n")
        print_with_color(:yellow, io, "Found in Iteration")
        @printf io ": "
        print_with_color(:bold, io, "$(n.iterations)\n")
        print_with_color(:blue, io, "Current Iteration")
        @printf io " : "
        print_with_color(:bold, io, "$(n.current_iteration)\n")
        print_with_color(:blue, io, "Technique")
        @printf io "         : "
        print_with_color(:bold, io, "$(n.technique)\n")
        print_with_color(:blue, io, "Function Calls")
        @printf io "    : "
        print_with_color(:bold, io, "$(n.cost_calls)\n")
        print_with_color(:blue, io, "  ***\n")
    end
    return
end


text:
export VarDelta

mutable struct VarDelta
    """
    Variable clamped to an observation
    """

    # Node properties
    id::Symbol
    marginal::Delta
    time::Union{Integer,Nothing}
    block::Bool

    # Keep track of prediction error
    pred_error::Float64
	Δpred_error::Float64

    function VarDelta(id::Symbol,
					  observed_value::Float64;
					  time::Union{Integer,Nothing}=0,
					  block=false)

        # Cast observed value to Delta distrbution
        marginal = Delta(observed_value)

		# Initialize prediction error
		pred_error = Inf
		Δpred_error = Inf

        return new(id, marginal, time, block, pred_error, Δpred_error)
    end
end

function entropy(variable::VarDelta)
    "An observed variable has no entropy."
    return 0.0
end

function gradient_entropy(variable::VarDelta)
    "Gradient of entropy for observed variable is 0."
    return 0.0
end

function prediction_error!(graph::MetaGraph, variable::VarDelta)
    "Compute prediction error of observation"

	# Reset prediction error
	pred_error_old = variable.pred_error

	# Find edge connected to node
	edge = edges(graph, graph[variable.id, :id])[1]

	# Extract message from factor to variable
	message_f2v = get_prop(graph, edge, :message_factor2var)

	# Compute and add prediction error
	variable.pred_error = -logpdf(message_f2v, mean(variable.marginal))

	# Compute change in prediction error
	variable.Δpred_error = norm(pred_error_old - variable.pred_error)
end

function belief(variable::VarDelta)
    "Edge is distribution"
    return Delta(variable.observation)
end

function act!(graph::MetaGraph, variable::VarDelta)
    "Pass message to edge with likelihood node"

    # Check each connected edge
    for edge in edges(graph, graph[variable.id, :id])

        # Check for incoming
		set_prop!(graph, edge, :message_var2factor, variable.marginal)
		set_prop!(graph, edge, :Δfree_energy, variable.Δpred_error)
    end
end

function react!(graph::MetaGraph, variable::VarDelta)
    "React to incoming messages"

    # Compute local free energy
    prediction_error!(graph, variable)

    # Message from edge to nodes
    act!(graph, variable)
end

max_stars_repo_name: wmkouw/schedulefreeMP

max_issues_repo_name: wmkouw/schedulefreeMP

max_forks_repo_name: wmkouw/schedulefreeMP


text:
import ..Inputs:
    allnamelists,
    allcards,
    required_namelists,
    optional_namelists,
    required_cards,
    optional_cards

export PWInput,
    exitfile,
    mkexitfile,
    allnamelists,
    allcards,
    required_namelists,
    optional_namelists,
    required_cards,
    optional_cards

"""
    PWInput <: QuantumESPRESSOInput
    PWInput(control, system, electrons, ions, cell, atomic_species, atomic_positions, k_points, cell_parameters)

Construct a `PWInput` which represents the input of program `pw.x`.

# Arguments
- `control::ControlNamelist=ControlNamelist()`: the `CONTROL` namelist of the input. Optional.
- `system::SystemNamelist=SystemNamelist()`: the `SYSTEM` namelist of the input. Optional.
- `electrons::ElectronsNamelist=ElectronsNamelist()`: the `ELECTRONS` namelist of the input. Optional.
- `ions::IonsNamelist=IonsNamelist()`: the `IONS` namelist of the input. Optional.
- `cell::CellNamelist=CellNamelist()`: the `CELL` namelist of the input. Optional.
- `atomic_species::AtomicSpeciesCard`: the `ATOMIC_SPECIES` card of the input. Must be provided explicitly.
- `atomic_positions::AtomicPositionsCard`: the `ATOMIC_POSITIONS` card of the input. Must be provided explicitly.
- `k_points::AbstractKPointsCard`: the `K_POINTS` card of the input. Must be provided explicitly.
- `cell_parameters::Union{Nothing,CellParametersCard}`: the `CELL_PARAMETERS` card of the input. Must be either `nothing` or a `CellParametersCard`.
"""
struct PWInput <: QuantumESPRESSOInput
    control::ControlNamelist
    system::SystemNamelist
    electrons::ElectronsNamelist
    ions::IonsNamelist
    cell::CellNamelist
    atomic_species::AtomicSpeciesCard
    atomic_positions::AtomicPositionsCard
    k_points::KPointsCard
    cell_parameters::Union{Nothing,CellParametersCard}
    constraints::Union{Union{Nothing,Float64}}
    occupations::Union{Nothing,Float64}
    atomic_forces::Union{Nothing,AtomicForcesCard}
end # struct PWInput
function PWInput(;
    control = ControlNamelist(),
    system,
    electrons = ElectronsNamelist(),
    ions = IonsNamelist(),
    cell = CellNamelist(),
    atomic_species,
    atomic_positions,
    k_points,
    cell_parameters = nothing,
    constraints = nothing,
    occupations = nothing,
    atomic_forces = nothing,
)
    @assert !isnothing(cell_parameters) || system.ibrav != 0 "`cell_parameters` is empty with `ibrav = 0`!"
    foreach(atomic_species.data) do datum
        path = joinpath(control.pseudo_dir, datum.pseudopot)
        if !isfile(path)
            @warn "pseudopotential file \"$path\" does not exist!"
        end
    end
    return PWInput(
        control,
        system,
        electrons,
        ions,
        cell,
        atomic_species,
        atomic_positions,
        k_points,
        cell_parameters,
        constraints,
        occupations,
        atomic_forces,
    )
end

exitfile(template::PWInput) = abspath(
    expanduser(joinpath(template.control.outdir, template.control.prefix * ".EXIT")),
)
function mkexitfile(template::PWInput)
    path = exitfile(template)
    mkpath(dirname(path))
    return touch(path)
end

"""
    allnamelists(input::PWInput)

Return an iterator of all `Namelist`s from a `PWInput`. You may want to `collect` them.
"""
allnamelists(input::PWInput) =
    (getfield(input, f) for f in (:control, :system, :electrons, :ions, :cell))

"""
    allcards(input::PWInput)

Get all `Card`s from a `PWInput`.
"""
allcards(input::PWInput) = (
    getfield(input, f) for f in (
        :atomic_species,
        :atomic_positions,
        :k_points,
        :cell_parameters,
        :constraints,
        :occupations,
        :atomic_forces,
    )
)

"""
    required_namelists(input::PWInput)

Return an iterator of required `Namelist`s from a `PWInput`. You may want to `collect` them.
"""
required_namelists(input::PWInput) = (getfield(input, f) for f in (:control, :system, :electrons))

"""
    optional_namelists(input::PWInput)

Return an iterator of optional `Namelist`s from a `PWInput`. You may want to `collect` them.
"""
optional_namelists(input::PWInput) = (getfield(input, f) for f in (:ions, :cell))

"""
    required_cards(input::PWInput)

Return an iterator of required `Card`s from a `PWInput`. You may want to `collect` them.
"""
required_cards(input::PWInput) =
    (getfield(input, f) for f in (:atomic_species, :atomic_positions, :k_points))

"""
    optional_cards(input::PWInput)

Return an iterator of optional `Card`s from a `PWInput`. You may want to `collect` them.
"""
optional_cards(input::PWInput) =
    (getfield(input, f) for f in (:cell_parameters, :constraints, :occupations, :atomic_forces))

"""
    getpotentials(input::PWInput)

Get the pseudopotential names from a `PWInput`.
"""
getpotentials(input::PWInput) = getpotentials(input.atomic_species)

"""
    getpseudodir(input::PWInput)

Get the directory storing the pseudopotential files.
"""
getpseudodir(input::PWInput) = getpseudodir(input.control)


text:
# default generating vector files
const K_3600_32_file = joinpath(@__DIR__(), "..", "generating_vectors", "K_3600_32.txt")
const CKN_250_20_file = joinpath(@__DIR__(), "..", "generating_vectors", "CKN_250_20.txt")

# re-compile if these files change
include_dependency(K_3600_32_file)
include_dependency(CKN_250_20_file)

# default generating vectors
const K_3600_32 = read32(K_3600_32_file)
const CKN_250_20 = read32(CKN_250_20_file)


text:
# Function for calculating ZV-MCMC estimators using linear polynomial
function linearZv(mcmc::Array{Float64, 2}, z::Array{Float64, 2})
  nPars = size(mcmc, 2)

  covAll = Array(Float64, nPars+1, nPars+1, nPars)
  precision = Array(Float64, nPars, nPars, nPars)
  sigma = Array(Float64, nPars, nPars)
  a = Array(Float64, nPars, nPars)

  for i = 1:nPars
     covAll[:, :, i] = cov([z mcmc[:, i]])
     precision[:, :, i] = inv(covAll[1:nPars, 1:nPars, i])
     sigma[:, i] = covAll[1:nPars, nPars+1, i]
     a[:, i] = -precision[:, :, i]*sigma[:, i]
  end

  zvMcmc = mcmc+z*a

  return zvMcmc, a
end

# Function for calculating ZV-MCMC estimators using quadratic polynomial
function quadraticZv(mcmc::Array{Float64, 2}, z::Array{Float64, 2})
  nData, nPars = size(mcmc)
  k = convert(Int, nPars*(nPars+3)/2)
  l = 2*nPars+1

  zQuadratic = Array(Float64, nData, k)
  covAll = Array(Float64, k+1, k+1, nPars)
  precision = Array(Float64, k, k, nPars)
  sigma = Array(Float64, k, nPars)
  a = Array(Float64, k, nPars)

  zQuadratic[:, 1:nPars] = z
  zQuadratic[:, (nPars+1):(2*nPars)] = 2*z.*mcmc-1
  for i = 1:(nPars-1)
    for j = (i+1):nPars
      zQuadratic[:, l] = mcmc[:, i].*z[:, j]+mcmc[:, j].*z[:, i]
      l += 1
    end
  end

  for i = 1:nPars
    covAll[:, :, i] = cov([zQuadratic mcmc[:, i]]);
    precision[:, :, i] =
      inv(covAll[1:k, 1:k, i])
    sigma[:, i] = covAll[1:k, k+1, i]
    a[:, i] = -precision[:, :, i]*sigma[:, i]
  end

  zvMcmc = mcmc+zQuadratic*a;

  return zvMcmc, a
end


text:
function insertsort!(A)
    for k in eachindex(A)
        while k>1 && A[k-1] > A[k]
            A[k-1], A[k] = A[k], A[k-1]
            k -= 1
        end
    end
    return A
end

function insertsortperm!(A)
    ind=collect(1:length(A))
    for k in eachindex(A)
        while k>1 && A[k-1] > A[k]
            A[k-1], A[k] = A[k], A[k-1]
            ind[k-1], ind[k] = ind[k], ind[k-1]
            k -= 1
        end
    end
    return ind
end

insertsort(A) = insertsort!(copy(A))

insertsortperm(A) = insertsortperm!(copy(A))


text:
using PyCall  # import packages
using Pkg
using Distributions, Plots
using DataFrames
using CSV
#load data
df1 = CSV.read("/home/lucifer/Downloads/archive (4)/diabetes.csv", DataFrame)
df = copy(df1)
#Remove outliers
df1 = df1[df1.Pregnancies .<14, :]
df1 = df1[df1.Glucose .> 50, :]
df1 = df1[df1.BloodPressure .< 120, :]
df1 = df1[df1.SkinThickness .< 80, :]
df1 = df1[df1.Insulin .< 600, :]
df1 = df1[df1.BMI .< 55, :]
df1 = df1[df1.DiabetesPedigreeFunction .< 2, :]
df1 = df1[df1.Age .< 70, :]

# Remove Skewness by passing the column values through the log function
Ins=[]
function remove_skew(x)
    for i in x
        if i>0||i!=0
            v = log(i)
            append!(Ins, v)
        else
            v = 0
            append!(Ins, v)
        end
    end
end

# Remove skewness  from the distribution
remove_skew(df1[!, :Insulin])
df1[!, :Insulin] = Float64.(Ins);
# feature scaling
using StatsBase
new_columns = names(df1)
n = length(names(df1))
c_names = names(df1)
for x in range(1, stop=n, length=n)
    x = Int64.(x)
    col = Float64.(df1[:, x])
    s = fit(ZScoreTransform, col)
    colm = (c_names[x]*"_"*(string(x)))
    df1[:, colm] = StatsBase.transform!(s, col)
end
# select features
features = select(df1, Between(:Pregnancies_1, :Age_8))

# select label
target = df1[!, :Outcome];

@pyimport sklearn.model_selection as ms
# split data into train|test data
X_train, X_test, y_train, y_test = ms.train_test_split(Array(features),Array(target), test_size = 0.2, random_state=42)

@pyimport sklearn
@pyimport xgboost as xgb
svm = pyimport("sklearn.svm") #for SVC
tree = pyimport("sklearn.tree") #for decision tree
neighbors = pyimport("sklearn.neighbors") #for kn
linear_model = pyimport("sklearn.linear_model") # for logistic regression
ensemble = pyimport("sklearn.ensemble") #for ensemble methods (ex. Randomforest,
                                        #AdaBoost, GradientBoosting)

                                        # list of models we are going to use for classification task
models = [xgb.XGBClassifier(),svm.SVC(),tree.DecisionTreeClassifier(),neighbors.KNeighborsClassifier(),linear_model.LogisticRegression(),ensemble.RandomForestClassifier(),ensemble.AdaBoostClassifier(),ensemble.GradientBoostingClassifier()]

metrics = pyimport("sklearn.metrics")
for model in models
    #fit train & test data to every model and print the accuracy
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)
    print(model, metrics.accuracy_score(y_test,y_pred),"\n")
end


text:
using Test, Suppressor, DataFrames, HMCResearchRandomWalks.PPD

#getdata_ppdstepsize([0.01, 0.05, 0.1])

@testset "setters for simulation control" begin
    set_PPD_ON(true)
    @test PPD.PPD_ON
    @test PPD.PPD_STEP_SIZE == 0.001
    @test PPD.STEP_SIZE_DICT["ppd"] == 0.001

    set_PPD_ON(false)
    @test !PPD.PPD_ON
    @test PPD.PPD_STEP_SIZE == 0.1
    @test PPD.STEP_SIZE_DICT["ppd"] == 0.1
end

"extra test, commented out for runtime"
# @testset "default ppd_sim" begin
#     set_PPD_ON()
#     set_NUMBER_OF_WALKS()
#     set_MAX_STEPS_PER_WALK()
#     output = @capture_out begin ppd_sim(8925) end
#     correct_output = string(
#     "############################\n",
#     "       PPD Experiment       \n",
#     "############################\n",
#     "_________Parameters_________\n",
#     "PPD on sensor:\t\ttrue\n",
#     "# of trials:\t\t1000\n",
#     "# of steps (max):\t1485604\n",
#     "Step size, water:\t0.1\n",
#     "Step size, PPD:\t\t0.001\n",
#     "Random seed:\t\t8925\n",
#     "_________Results____________\n",
#     "# in sensor:\t\t462\n",
#     "# of escaped:\t\t538\n",
#     "# unresolved:\t\t0\n"
#     )
#     @test output === correct_output
# end

@testset "fast custom ppd_sim" begin
    set_PPD_ON(false)
    set_NUMBER_OF_WALKS(10)
    set_MAX_STEPS_PER_WALK(1000000)
    output = @capture_out begin ppd_sim(5534) end
    correct_output = string(
    "############################\n",
    "       PPD Experiment       \n",
    "############################\n",
    "_________Parameters_________\n",
    "PPD on sensor:\t\tfalse\n",
    "# of trials:\t\t10\n",
    "# of steps (max):\t1000000\n",
    "Step size, water:\t0.1\n",
    "Random seed:\t\t5534\n",
    "_________Results____________\n",
    "# in sensor:\t\t10\n",
    "# of escaped:\t\t0\n",
    "# unresolved:\t\t0\n"
    )
    @test output === correct_output
end

@testset "ppd_print" begin
    n = 420
    m = 2000
    seed = 1234
    set_NUMBER_OF_WALKS(n)
    set_MAX_STEPS_PER_WALK(m)
    output = @capture_out begin PPD.ppd_print(seed) end
    correct_output = string(
    "############################\n",
    "   Compare PPD Step Sizes   \n",
    "############################\n",
    "_________Parameters_________\n",
    "# of trials:\t\t$n\n",
    "# of steps (max):\t$m\n",
    "Step size, water:\t0.1\n",
    "Random seed:\t\t$seed\n"
    )
    @test output === correct_output
end

@testset "getdata_ppd..." begin
    set_NUMBER_OF_WALKS(50)
    set_MAX_STEPS_PER_WALK()
    @suppress begin
        global getdata_ppd_test_df = getdata_ppdstepsize([0.001, 0.002], 6520)
    end
    correct_df = DataFrame(
        ppd_step_size = [0.001, 0.002],
        sensor_yield = [20, 38],
        escaped = [30, 12],
        unresolved = [0, 0]
    )
    @test getdata_ppd_test_df == correct_df
end

max_stars_repo_name: s-y-wu/HMCResearchRandomWalks

max_issues_repo_name: s-y-wu/HMCResearchRandomWalks

max_forks_repo_name: s-y-wu/HMCResearchRandomWalks


text:
module GrammaticalEvolution

import Base.*
import Base.sort!
import Base.getindex
import Base.setindex!
import Base.length
import Base.endof
import Base.isless
import Base.push!
import Base.pop!

export Individual, Population
export select_two_individuals, one_point_crossover, mutate!, evaluate!, generate, transform
export length, getindex, endof, setindex!, isless, genome_iterator
export MaxWrapException

include("EBNF.jl")

export Grammar, @grammar, Rule, AndRule, OrRule, parseGrammar

type MaxWrapException <: Exception end

abstract Individual
abstract Population

# methods that have to be supported by subclasses of population
length{T <: Population}(pop::T) = length(pop.individuals)
getindex{T <: Population}(pop::T, indices...) = pop.individuals[indices...]
push!{T <: Population, S <: Individual}(pop::T, ind::S) = push!(pop.individuals, ind)
pop!{T <: Population}(pop::T) = pop!(pop.individuals)

# methods that have to be supported by subclasses of individuals
length{T <: Individual}(ind::T) = length(ind.genome)
endof{T <: Individual}(ind::T) = endof(ind.genome)
getindex{T <: Individual}(ind::T, indices...) = ind.genome[indices...]
setindex!{T <: Individual}(ind::T, value::Int64, indices) = ind.genome[indices] = value
isless{T <: Individual}(ind1::T, ind2::T) = ind1.fitness < ind2.fitness
getFitness{T <: Individual}(ind::T) = ind.fitness
# evaluate(ind::Individual) = nothing
evaluate!{T <: Individual}(grammar::Grammar, ind::T, args...) = nothing

# TODO: this should be distributed
function evaluate!{PopulationType <: Population}(grammar::Grammar, pop::PopulationType, args...)
  for i=1:length(pop)
    if getFitness(pop[i]) < 0.0
      evaluate!(grammar, pop[i], args...)
    end
  end
end

function sort!{PopulationType <: Population}(pop::PopulationType)
  sort!(pop.individuals)
end

function select_two_individuals{T <: Individual}(individuals::Array{T,1})
  # randomly select two individuals from a population
  while true
    i1 = rand(1:length(individuals))
    i2 = rand(1:length(individuals))

    # make sure the two individuals are not the same
    if i1 != i2
      return (i1, i2)
    end
  end
end

function one_point_crossover{IndividualType <: Individual}(ind1::IndividualType, ind2::IndividualType)
  cross_point = rand(1:length(ind1))
  g1 = vcat(ind1[1:cross_point-1], ind2[cross_point:end])
  g2 = vcat(ind2[1:cross_point-1], ind1[cross_point:end])

  return (IndividualType(g1), IndividualType(g2), cross_point)
end

function mutate!(ind::Individual, mutation_rate::Float64; max_value=1000)
  for i=1:length(ind)
    if rand() < mutation_rate
      ind[i] = rand(1:max_value)
    end
  end
end

function generate{PopulationType <: Population}(grammar::Grammar, population::PopulationType, top_percent::Float64, prob_mutation::Float64, mutation_rate::Float64, args...)
  # sort population
  sort!(population)

  # take top performers for cross-over
  top_num::Int64 = floor(length(population)*top_percent)
  top_performers = population[1:top_num]

  # create a new population
  genome_size = length(population[1])
  new_population = PopulationType(top_performers)

  # re-populate by mating top performers
  while length(new_population) < length(population)
    (i1, i2) = select_two_individuals(top_performers)
    (ind1, ind2) = one_point_crossover(top_performers[i1], top_performers[i2])
    push!(new_population, ind1)
    push!(new_population, ind2)
  end

  # mutate newly created population
  for j=(top_num+1):length(population)
    if rand() < prob_mutation
      mutate!(new_population[j], mutation_rate)
    end
  end

  # evaluate new population and re-sort
  evaluate!(grammar, new_population, args...)
  sort!(new_population)

  # it's possible that we might have added too many individuals, so trim down if necessary
  while length(new_population) > length(population)
    pop!(new_population)
  end

  return new_population
end

# stateful iterator that keeps track of its current position, wraps the position
# when the maximum length is reached, and emits an exception when the maximum
# number of wraps occurs
function genome_iterator(size::Int64, maxwraps::Int64)
  i::Int64 = 1
  wraps::Int64 = 0

  function next()
    while true
      produce(i)

      i += 1
      if i > size
        wraps += 1
        i = 1
      end

      if wraps > maxwraps
        throw(MaxWrapException())
      end
    end
  end

  return Task(next)
end

function transform(grammar::Grammar, ind::Individual; maxwraps=2)
  pos = genome_iterator(length(ind), maxwraps)
  value = transform(grammar, grammar.rules[:start], ind, pos)
  return value
end

function transform(grammar::Grammar, rule::OrRule, ind::Individual, pos::Task)
  idx = (ind[consume(pos)] % length(rule.values))+1
  value = transform(grammar, rule.values[idx], ind, pos)

  if rule.action !== nothing
    value = rule.action(value)
  end

  return value
end

function transform(grammar::Grammar, rule::RangeRule, ind::Individual, pos::Task)
  value = (ind[consume(pos)] % length(rule.range))+rule.range.start

  if rule.action !== nothing
    value = rule.action(value)
  end

  return value
end

function transform(grammar::Grammar, rule::ReferencedRule, ind::Individual, pos::Task)
  return transform(grammar, grammar.rules[rule.symbol], ind, pos)
end

function transform(grammar::Grammar, rule::Terminal, ind::Individual, pos::Task)
  return rule.value
end

function transform(grammar::Grammar, rule::AndRule, ind::Individual, pos::Task)
  values = [transform(grammar, subrule, ind, pos) for subrule in rule.values]

  if rule.action !== nothing
    values = rule.action(values)
  end

  return values
end

function transform(grammar::Grammar, sym::Symbol, ind::Individual, pos::Task)
  return sym
end

function transform(grammar::Grammar, q::QuoteNode, ind::Individual, pos::Task)
  return q.value
end

function transform(grammar::Grammar, rule::ExprRule, ind::Individual, pos::Task)
  args = [transform(grammar, arg, ind, pos) for arg in rule.args]
  return Expr(args...)
end

# It's very unlikely these two methods will be useful -- the maximum size of the genome is arbritrarily high, so
# you'll likely end up with mostly large numbers
# function transform(grammar::Grammar, rule::ZeroOrMoreRule, ind::Individual, pos::Task)
#   # genome value gives number of time to repeat
#   reps = ind[consume(pos)]

#   # invoke given rule reps times
#   values = [transform(grammar, rule.rule, ind, pos) for i=1:reps]

#   if rule.action !== nothing
#     values = rule.action(values)
#   end

#   return values
# end

# function transform(grammar::Grammar, rule::OneOrMoreRule, ind::Individual, pos::Task)
#   # genome value gives number of time to repeat
#   reps = ind[consume(pos)]

#   # enforce that it's at least one
#   if reps == 0
#     reps = 1
#   end

#   # invoke given rule reps times
#   values = [transform(grammar, rule.value, ind, pos) for i=1:reps]

#   if rule.action !== nothing
#     values = rule.action(values)
#   end

#   return values
# end

function transform(grammar::Grammar, rule::RepeatedRule, ind::Individual, pos::Task)
  # genome value gives number of time to repeat
  reps = ind[consume(pos)]
  range = (rule.range.stop - rule.range.start)
  reps = (reps % range) + rule.range.start

  # invoke given rule reps times
  values = [transform(grammar, rule.value, ind, pos) for i=1:reps]

  if rule.action !== nothing
    values = rule.action(values)
  end

  return values
end

end # module

max_stars_repo_name: abeschneider/GrammaticalEvolution

max_issues_repo_name: abeschneider/GrammaticalEvolution

max_forks_repo_name: abeschneider/GrammaticalEvolution


text:
"""
$(TYPEDEF)

Real vector cone of squares.
"""
struct VectorCSqr{T <: Real} <: ConeOfSquares{T} end

"""
$(TYPEDSIGNATURES)

The rank of the vector cone of squares, equal to the vector length.
"""
vector_dim(::Type{<:VectorCSqr}, d::Int) = d

mutable struct VectorCSqrCache{T <: Real} <: CSqrCache{T}
    viw::Vector{T}
    wi::Vector{T}
    ϕ::T
    ζ::T
    ζi::T
    ζivi::T
    σ::T
    ∇h::Vector{T}
    ∇2h::Vector{T}
    ∇3h::Vector{T}

    m::Vector{T}
    α::Vector{T}
    γ::Vector{T}
    k1::T
    k2::T
    k3::T

    w1::Vector{T}
    w2::Vector{T}

    VectorCSqrCache{T}() where {T <: Real} = new{T}()
end

function setup_csqr_cache(cone::EpiPerSepSpectral{VectorCSqr{T}}) where T
    cone.cache = cache = VectorCSqrCache{T}()
    d = cone.d
    cache.viw = zeros(T, d)
    cache.wi = zeros(T, d)
    cache.∇h = zeros(T, d)
    cache.∇2h = zeros(T, d)
    cache.∇3h = zeros(T, d)
    cache.w1 = zeros(T, d)
    cache.w2 = zeros(T, d)
    cache.m = zeros(T, d)
    cache.α = zeros(T, d)
    cache.γ = zeros(T, d)
    return
end

function set_initial_point!(
    arr::AbstractVector,
    cone::EpiPerSepSpectral{<:VectorCSqr},
    )
    (arr[1], arr[2], w0) = get_initial_point(cone.d, cone.h)
    @views fill!(arr[3:end], w0)
    return arr
end

function update_feas(cone::EpiPerSepSpectral{VectorCSqr{T}}) where T
    @assert !cone.feas_updated
    cache = cone.cache
    v = cone.point[2]

    if (v > eps(T)) && all(>(eps(T)), cone.w_view)
        @. cache.viw = cone.w_view / v
        cache.ϕ = h_val(cache.viw, cone.h)
        cache.ζ = cone.point[1] - v * cache.ϕ
        cone.is_feas = (cache.ζ > eps(T))
    else
        cone.is_feas = false
    end

    cone.feas_updated = true
    return cone.is_feas
end

function is_dual_feas(cone::EpiPerSepSpectral{VectorCSqr{T}}) where T
    u = cone.dual_point[1]
    (u < eps(T)) && return false
    @views w = cone.dual_point[3:end]

    h_conj_dom_pos(cone.h) && any(<(eps(T)), w) && return false
    uiw = cone.cache.w1
    @. uiw = w / u
    return (cone.dual_point[2] - u * h_conj(uiw, cone.h) > eps(T))
end

function update_grad(cone::EpiPerSepSpectral{<:VectorCSqr})
    @assert !cone.grad_updated && cone.is_feas
    grad = cone.grad
    v = cone.point[2]
    cache = cone.cache
    ζi = cache.ζi = inv(cache.ζ)
    ∇h = cache.∇h
    h_der1(∇h, cache.viw, cone.h)
    cache.σ = cache.ϕ - dot(cache.viw, ∇h)
    @. cache.wi = inv(cone.w_view)

    grad[1] = -ζi
    grad[2] = -inv(v) + ζi * cache.σ
    @. grad[3:end] = -cache.wi + ζi * ∇h

    cone.grad_updated = true
    return grad
end

function update_hess_aux(cone::EpiPerSepSpectral{<:VectorCSqr})
    @assert !cone.hess_aux_updated
    @assert cone.grad_updated
    cache = cone.cache
    cache.ζivi = inv(cache.ζ * cone.point[2])
    h_der2(cache.∇2h, cache.viw, cone.h)
    cone.hess_aux_updated = true
end

function update_hess(cone::EpiPerSepSpectral{<:VectorCSqr})
    cone.hess_aux_updated || update_hess_aux(cone)
    isdefined(cone, :hess) || alloc_hess!(cone)
    v = cone.point[2]
    H = cone.hess.data
    cache = cone.cache
    ζi = cache.ζi
    ζivi = cache.ζivi
    σ = cache.σ
    viw = cache.viw
    wi = cache.wi
    ∇h = cache.∇h
    ∇2h = cache.∇2h

    # Huu
    Huu = H[1, 1] = abs2(ζi)

    # Huv
    H[1, 2] = -Huu * σ

    # Hvv start
    Hvv = v^-2 + abs2(ζi * σ)

    @inbounds for j in 1:cone.d
        ζi∇h_j = ζi * ∇h[j]
        ζi2∇h_j = ζi * ζi∇h_j
        ζivi∇2h_j = ζivi * ∇2h[j]
        term_j = viw[j] * ζivi∇2h_j
        Hvv += viw[j] * term_j
        j2 = 2 + j

        # Huw
        Huwj = H[1, j2] = -ζi2∇h_j

        # Hvw
        H[2, j2] = σ * ζi2∇h_j - term_j

        # Hww
        for i in 1:(j - 1)
            H[2 + i, j2] = ζi2∇h_j * ∇h[i]
        end
        H[j2, j2] = abs2(ζi∇h_j) + ζivi∇2h_j + abs2(wi[j])
    end

    # Hvv end
    H[2, 2] = Hvv

    cone.hess_updated = true
    return cone.hess
end

function hess_prod!(
    prod::AbstractVecOrMat{T},
    arr::AbstractVecOrMat{T},
    cone::EpiPerSepSpectral{VectorCSqr{T}},
    ) where T
    cone.hess_aux_updated || update_hess_aux(cone)
    v = cone.point[2]
    w = cone.w_view
    cache = cone.cache
    ζi = cache.ζi
    ζivi = cache.ζivi
    viw = cache.viw
    wi = cache.wi
    σ = cache.σ
    ∇h = cache.∇h
    ∇2h = cache.∇2h
    ξb = cache.w1

    @inbounds for j in 1:size(arr, 2)
        p = arr[1, j]
        q = arr[2, j]
        @views r = arr[3:end, j]

        viq = q / v
        @. ξb = ζivi * ∇2h * (r - viq * w)
        c1 = -ζi * (p - σ * q - dot(∇h, r)) * ζi

        prod[1, j] = -c1
        prod[2, j] = c1 * σ - dot(viw, ξb) + viq / v
        @. prod[3:end, j] = c1 * ∇h + ξb + wi * r * wi
    end

    return prod
end

function update_inv_hess_aux(cone::EpiPerSepSpectral{<:VectorCSqr})
    @assert !cone.inv_hess_aux_updated
    cone.hess_aux_updated || update_hess_aux(cone)
    v = cone.point[2]
    cache = cone.cache
    ∇h = cache.∇h
    wi = cache.wi
    w1 = cache.w1
    m = cache.m
    α = cache.α
    γ = cache.γ

    @. w1 = cache.ζivi * cache.∇2h
    @. m = inv(w1 + abs2(wi))
    @. α = m * ∇h
    @. γ = m * cache.viw * w1

    cache.k1 = abs2(cache.ζ) + dot(∇h, α)
    cache.k2 = cache.σ + dot(∇h, γ)
    cache.k3 = (inv(v) + dot(wi, γ)) / v

    cone.inv_hess_aux_updated = true
end

function update_inv_hess(cone::EpiPerSepSpectral{<:VectorCSqr})
    cone.inv_hess_aux_updated || update_inv_hess_aux(cone)
    isdefined(cone, :inv_hess) || alloc_inv_hess!(cone)
    Hi = cone.inv_hess.data
    cache = cone.cache
    m = cache.m
    α = cache.α
    γ = cache.γ
    k2 = cache.k2
    k3 = cache.k3

    # Hiuu, Hiuv, Hivv
    k23i = k2 / k3
    Hi[1, 1] = cache.k1 + k23i * k2
    Hi[1, 2] = k23i
    Hi[2, 2] = inv(k3)

    @inbounds for j in 1:cone.d
        j2 = 2 + j

        # Hivw
        Hivj = Hi[2, j2] = γ[j] / k3

        # Hiuw
        Hi[1, j2] = α[j] + k2 * Hivj

        # Hiww
        for i in 1:j
            Hi[2 + i, j2] = Hivj * γ[i]
        end
        Hi[j2, j2] += m[j]
    end

    cone.inv_hess_updated = true
    return cone.inv_hess
end

function inv_hess_prod!(
    prod::AbstractVecOrMat{T},
    arr::AbstractVecOrMat{T},
    cone::EpiPerSepSpectral{VectorCSqr{T}},
    ) where T
    cone.inv_hess_aux_updated || update_inv_hess_aux(cone)
    cache = cone.cache
    m = cache.m
    α = cache.α
    γ = cache.γ
    k1 = cache.k1
    k2 = cache.k2
    k3 = cache.k3

    @inbounds for j in 1:size(arr, 2)
        p = arr[1, j]
        @views r = arr[3:end, j]

        cv = (k2 * p + arr[2, j] + dot(γ, r)) / k3

        prod[1, j] = k1 * p + k2 * cv + dot(α, r)
        prod[2, j] = cv
        @. @views prod[3:end, j] = p * α + cv * γ + m * r
    end

    return prod
end

function update_dder3_aux(cone::EpiPerSepSpectral{<:VectorCSqr})
    @assert !cone.dder3_aux_updated
    cone.hess_aux_updated || update_hess_aux(cone)
    h_der3(cone.cache.∇3h, cone.cache.viw, cone.h)
    cone.dder3_aux_updated = true
end

function dder3(
    cone::EpiPerSepSpectral{VectorCSqr{T}},
    dir::AbstractVector{T},
    ) where T
    cone.dder3_aux_updated || update_dder3_aux(cone)
    v = cone.point[2]
    w = cone.w_view
    dder3 = cone.dder3
    cache = cone.cache
    ζi = cache.ζi
    σ = cache.σ
    ∇h = cache.∇h
    ∇2h = cache.∇2h
    wi = cache.wi
    ξ = cache.w1
    ξb = cache.w2

    p = dir[1]
    q = dir[2]
    @views r = dir[3:end]

    viq = q / v
    @. ξ = r - viq * w
    @. ξb = cache.ζivi * ∇2h * ξ
    ζiχ = ζi * (p - σ * q - dot(∇h, r))
    ξbξ = dot(ξb, ξ) / 2
    c1 = -ζi * (ζiχ^2 + ξbξ)

    c2 = -ζi / 2
    ξ ./= v
    w_aux = ξb
    w_aux .*= ζiχ + viq
    @. w_aux += c2 * ξ * cache.∇3h * ξ

    dder3[1] = -c1
    dder3[2] = c1 * σ - dot(cache.viw, w_aux) + (ξbξ + viq^2) / v
    @. dder3[3:end] = c1 * ∇h + w_aux + abs2(r * wi) * wi

    return dder3
end


text:
# DBitArray is postponed after looking into internals. There is a chance of data race.
# For now without concern for space complexity lets just use regular boolean value with 8-bit footprint
# We can optimize it once we know BitArrays can be used safely or with minimal work.
# Lets focus on accuracy first.
#

mutable struct DBool
    bit::Bool
    DBool(b::Bool) = new(b)
end

mutable struct DBitArray
    bitarray::BitArray
    function DBitArray(count::Int)
        return new(bitrand(count))
    end
end

@forward DBitArray.bitarray Base.getindex, Base.setindex!

function set!(b::DBool, x::Bool)
    b.bit = x
    b
end

(b::DBool)(x::Bool) = set!(b, x)

Base.xor(a::DBool, b::DBool) = let x = Base.xor(a.bit, b.bit); DBool(x) end
and(a::DBool, b::DBool) = let x = Base.and_int(a.bit, b.bit); DBool(x) end
or(a::DBool, b::DBool) = let x = Base.or_int(a.bit, b.bit); DBool(x) end
not(a::DBool) = let x = Base.not_int(a.bit); DBool(x) end

(&)(a::DBool, b::DBool) = and(a, b)
(|)(a::DBool, b::DBool) = or(a, b)
(!)(a::DBool) = not(a::DBool)
(==)(a::DBool, b::DBool) = DBool(a.bit == b.bit)
(⊻)(a::DBool, b::DBool) = xor(a, b)
# Refer to bool.jl for other implementations

@adjoint DBool(b::Bool) = DBool(b), (b) -> (DBool(b), )

@adjoint xor(a::DBool, b::DBool) = xor(a, b), (∂) -> (∂, ∂) # TODO
@adjoint and(a::DBool, b::DBool) = and(a, b), (∂) -> (!xor(∂, b), !xor(∂, a))
@adjoint or(a::DBool, b::DBool) = or(a, b), (∂) -> (!xor(∂, b), !xor(∂, a)) # TODO
@adjoint not(a::DBool) = not(a), (∂) -> (∂,)

# @adjoint (&)(a::DBool, b::DBool) = and(a, b), (∂) -> (xor(∂, b), xor(∂, a))
# @adjoint (|)(a::DBool, b::DBool) = or(a, b), (∂) -> (!xor(∂, b), !xor(∂, a))
# @adjoint (!)(a::DBool) = not(a), (∂) -> (∂ == a,)

⨁(a::DBool, b::DBool) = !xor(a, b) # U2A01-3 Remember

# TODO sessions are important here for others to work in this space

# Array versions
# Broadcasts works out of box
# Adjoints need to supported.

Base.xor(a::Array{DBool}, b::Array{DBool}) = xor.(a, b)
and(a::Array{DBool}, b::Array{DBool}) = and.(a, b)
or(a::Array{DBool}, b::Array{DBool}) = or.(a, b)
not(a::Array{DBool}) = not.(a)

(&)(a::Array{DBool}, b::Array{DBool}) = and(a, b)
(|)(a::Array{DBool}, b::Array{DBool}) = or(a, b)
(==)(a::Array{DBool}, b::Array{DBool}) = a .== b
(⊻)(a::Array{DBool}, b::Array{DBool}) = xor(a, b)

import Base: broadcasted

(!)(a::Array{DBool}) = not.(a)

@adjoint broadcasted(::typeof(xor), x::Array{DBool}, y::Array{DBool}) = xor.(x, y), Δ -> (nothing, Δ, Δ)
@adjoint broadcasted(::typeof(and), x::Array{DBool}, y::Array{DBool}) = and.(a, b), (∂) -> (nothing, !xor.(∂, b), !xor.(∂, a))
@adjoint broadcasted(::typeof(or), x::Array{DBool}, y::Array{DBool}) = or.(a, b), (∂) -> (nothing, !xor.(∂, b), !xor.(∂, a))
@adjoint broadcasted(::typeof(not), x::Array{DBool}) = not.(a), (∂) -> (nothing, ∂)
@adjoint broadcasted(::typeof(⨁), x::Array{DBool}, y::Array{DBool}) = xor.(x, y), Δ -> (nothing, !Δ, !Δ)

# Not sure if this broadcasted version
# @adjoint DBool(b::Bool) = DBool(b), (b) -> (DBool(b), )

# @adjoint (&)(a::DBool, b::DBool) = and(a, b), (∂) -> (xor(∂, b), xor(∂, a))
# @adjoint (|)(a::DBool, b::DBool) = or(a, b), (∂) -> (!xor(∂, b), !xor(∂, a))
# @adjoint (!)(a::DBool) = not(a), (∂) -> (∂ == a,)

lossF(x, y) = sum(xor(x, y))


text:
"""
    _make_gen_map(case)

Given a Case object, build a sparse matrix representing generator topology.
"""
function _make_gen_map(case::Case)::SparseMatrixCSC
    num_bus = length(case.busid)
    bus_idx = 1:num_bus
    bus_id2idx = Dict(case.busid .=> bus_idx)
    num_gen = length(case.genid)
    gen_idx = 1:num_gen
    gen_bus_idx = [bus_id2idx[b] for b in case.gen_bus]
    gen_map = sparse(gen_bus_idx, gen_idx, 1, num_bus, num_gen)
    return gen_map
end

"""
    _make_branch_map(case)

Given a Case object, build a sparse matrix representing branch topology.
"""
function _make_branch_map(case::Case)::SparseMatrixCSC
    num_branch_ac = length(case.branchid)
    num_branch = num_branch_ac + length(case.dclineid)
    num_bus = length(case.busid)
    branch_idx = 1:num_branch
    bus_idx = 1:num_bus
    bus_id2idx = Dict(case.busid .=> bus_idx)
    all_branch_to = vcat(case.branch_to, case.dcline_to)
    all_branch_from = vcat(case.branch_from, case.dcline_from)
    branch_to_idx = [bus_id2idx[b] for b in all_branch_to]
    branch_from_idx = [bus_id2idx[b] for b in all_branch_from]
    branches_to = sparse(branch_to_idx, branch_idx, 1, num_bus, num_branch)
    branches_from = sparse(branch_from_idx, branch_idx, -1, num_bus, num_branch)
    branch_map = branches_to + branches_from
    return branch_map
end

"""
    _make_bus_demand_weighting(case)

Given a Case object, build a sparse matrix that indicates the weighting of each bus in
    each zone.
"""
function _make_bus_demand_weighting(case::Case)::SparseMatrixCSC
    bus_idx = 1:length(case.busid)
    bus_df = DataFrames.DataFrame(;
        name=case.busid, load=case.bus_demand, zone=case.bus_zone
    )
    zone_demand = DataFrames.combine(DataFrames.groupby(bus_df, :zone), :load => sum)
    zone_list = sort(collect(Set(case.bus_zone)))
    zone_idx = 1:length(zone_list)
    zone_id2idx = Dict(zone_list .=> zone_idx)
    bus_df_with_zone_load = DataFrames.innerjoin(bus_df, zone_demand; on=:zone)
    bus_share = bus_df[:, :load] ./ bus_df_with_zone_load[:, :load_sum]
    bus_zone_idx = Int64[zone_id2idx[z] for z in case.bus_zone]
    zone_to_bus_shares = sparse(bus_zone_idx, bus_idx, bus_share)::SparseMatrixCSC
    return zone_to_bus_shares
end

"""
    _make_bus_demand(case)

Given a Case object, build a matrix of demand by (bus, hour) for this interval.
"""
function _make_bus_demand(case::Case, start_index::Int, end_index::Int)::Matrix
    # Bus weighting
    zone_to_bus_shares = _make_bus_demand_weighting(case)

    # Profiles
    simulation_demand = Matrix(case.demand[start_index:end_index, 2:end])
    bus_demand = permutedims(simulation_demand * zone_to_bus_shares)
    return bus_demand
end

"""
    _make_bus_demand_flexibility_amount(case, demand_flexibility)

Given a Case object and a DemandFlexibility object, build a matrix of demand flexibility
    by (bus, hour) for this interval.
"""
function _make_bus_demand_flexibility_amount(
    case::Case, demand_flexibility::DemandFlexibility, start_index::Int, end_index::Int
)::Tuple{Matrix,Matrix}
    # Bus weighting
    zone_to_bus_shares = _make_bus_demand_weighting(case)
    # Demand flexibility up profile
    simulation_demand_flex_amt_up = Matrix(
        demand_flexibility.flex_amt_up[start_index:end_index, 2:end]
    )
    # Demand flexibility down profile
    simulation_demand_flex_amt_dn = Matrix(
        demand_flexibility.flex_amt_dn[start_index:end_index, 2:end]
    )
    # convert per-area input to per-bus values
    if demand_flexibility.input_granularity == "AREA"
        bus_demand_flex_amt_up = permutedims(
            simulation_demand_flex_amt_up * zone_to_bus_shares
        )
        bus_demand_flex_amt_dn = permutedims(
            simulation_demand_flex_amt_dn * zone_to_bus_shares
        )
        # input is already per-bus, copy to new container
    elseif demand_flexibility.input_granularity == "BUS"
        bus_demand_flex_amt_up = permutedims(simulation_demand_flex_amt_up)
        bus_demand_flex_amt_dn = permutedims(simulation_demand_flex_amt_dn)
    else
        println("Warning! Invalid demand flexibility parameter input_granularity!")
    end
    return (bus_demand_flex_amt_up, bus_demand_flex_amt_dn)
end

"""
    _build_segment_slope(case, segment_idx, segment_width)

Given a Case object, an index of segments (e.g. 1:3), and a vector of widths
    (matching dimension of gen), return segment slopes by (gen, segment_idx).
"""
function _build_segment_slope(case::Case, segment_idx, segment_width)::Matrix
    # Note: this formulation still assumes quadratic cost curves only!
    COST = 5        # Positional index from mpc.gencost
    segment_slope = zeros(length(case.genid), length(segment_idx))
    for i in segment_idx
        segment_slope[:, i] = (
            (2 * case.gencost_orig[:, COST] .* case.gen_pmin) +
            case.gencost_orig[:, COST + 1] +
            (2 * i - 1) * case.gencost_orig[:, COST] .* segment_width
        )
    end
    return segment_slope
end

function _make_sets(
    case::Case;
    storage::Union{Storage,Nothing}=nothing,
    demand_flexibility::Union{DemandFlexibility,Nothing}=nothing,
)::Sets
    # Positional indices from mpc.gencost
    MODEL = 1
    NCOST = 4
    # Sets - Buses
    num_bus = length(case.busid)
    bus_idx = 1:num_bus
    bus_id2idx = Dict(case.busid .=> bus_idx)
    load_bus_idx = findall(case.bus_demand .> 0)
    num_load_bus = length(load_bus_idx)
    load_bus_map = sparse(load_bus_idx, 1:num_load_bus, 1, num_bus, num_load_bus)
    # Sets - branches
    ac_branch_rating = replace(case.branch_rating, 0 => Inf)
    branch_rating = vcat(ac_branch_rating, case.dcline_pmax)
    num_branch_ac = length(case.branchid)
    num_branch = num_branch_ac + length(case.dclineid)
    branch_idx = 1:num_branch
    noninf_branch_idx = findall(branch_rating .!= Inf)
    all_branch_to = vcat(case.branch_to, case.dcline_to)
    all_branch_from = vcat(case.branch_from, case.dcline_from)
    branch_to_idx = Int64[bus_id2idx[b] for b in all_branch_to]
    branch_from_idx = Int64[bus_id2idx[b] for b in all_branch_from]
    # Sets - generators
    num_gen = length(case.genid)
    gen_idx = 1:num_gen
    # Subsets - generators
    gen_wind_idx = gen_idx[findall(
        (case.genfuel .== "wind") .| (case.genfuel .== "wind_offshore")
    )]
    gen_solar_idx = gen_idx[findall(case.genfuel .== "solar")]
    gen_hydro_idx = gen_idx[findall(case.genfuel .== "hydro")]
    renewable_idx = sort(vcat(gen_wind_idx, gen_solar_idx, gen_hydro_idx))
    case.gen_pmax[renewable_idx] .= Inf
    noninf_pmax = findall(case.gen_pmax .!= Inf)
    num_wind = length(gen_wind_idx)
    num_solar = length(gen_solar_idx)
    num_hydro = length(gen_hydro_idx)
    # Generator cost curve segments
    piecewise_enabled = (sum(case.gencost[:, MODEL] .== 1) > 0)
    @assert(
        piecewise_enabled,
        "No piecewise segments detected. " * "Did you forget to linearize_gencost?"
    )
    num_segments = convert(Int, maximum(case.gencost[:, NCOST])) - 1
    segment_idx = 1:num_segments
    # Demand flexibility, additional info for input->bus conversion
    if demand_flexibility.enabled
        if demand_flexibility.input_granularity == "BUS"
            # assume each flexible bus can go up/dn, so only use the Dataframe for up
            flexible_bus_str = names(demand_flexibility.flex_amt_up)[2:end]
            flexible_bus_id = [parse(Int64, bus) for bus in flexible_bus_str]
            flexible_bus_idx = [bus_id2idx[bus] for bus in flexible_bus_id]
            num_flexible_bus = length(flexible_bus_idx)
            flexible_load_bus_map = sparse(
                flexible_bus_idx, 1:num_flexible_bus, 1, num_bus, num_flexible_bus
            )::SparseMatrixCSC
        elseif demand_flexibility.input_granularity == "AREA"
            flexible_bus_idx = load_bus_idx
            num_flexible_bus = num_load_bus
            flexible_load_bus_map = load_bus_map
        end
    else
        flexible_bus_idx = nothing
        num_flexible_bus = 0
        flexible_load_bus_map = nothing
    end
    # Storage
    storage_enabled = isa(storage, Storage) && storage.enabled
    num_storage = storage_enabled ? size(storage.gen, 1) : 0
    storage_idx = storage_enabled ? (1:num_storage) : nothing
    sets = Sets(;
        num_bus=num_bus,
        bus_idx=bus_idx,
        bus_id2idx=bus_id2idx,
        load_bus_idx=load_bus_idx,
        num_load_bus=num_load_bus,
        load_bus_map=load_bus_map,
        num_branch=num_branch,
        num_branch_ac=num_branch_ac,
        branch_idx=branch_idx,
        noninf_branch_idx=noninf_branch_idx,
        branch_to_idx=branch_to_idx,
        branch_from_idx=branch_from_idx,
        num_gen=num_gen,
        gen_idx=gen_idx,
        noninf_pmax=noninf_pmax,
        gen_hydro_idx=gen_hydro_idx,
        gen_solar_idx=gen_solar_idx,
        gen_wind_idx=gen_wind_idx,
        renewable_idx=renewable_idx,
        num_wind=num_wind,
        num_solar=num_solar,
        num_hydro=num_hydro,
        num_segments=num_segments,
        segment_idx=segment_idx,
        num_storage=num_storage,
        storage_idx=storage_idx,
        flexible_bus_idx=flexible_bus_idx,
        num_flexible_bus=num_flexible_bus,
        flexible_load_bus_map=flexible_load_bus_map,
    )
    return sets
end

function _add_constraint_power_balance!(
    m::JuMP.Model,
    case::Case,
    sets::Sets,
    storage::Storage,
    demand_flexibility::DemandFlexibility,
    bus_demand::Matrix,
    load_shed_enabled::Bool,
)
    # Generator topology matrix
    gen_map = _make_gen_map(case)
    # Branch connectivity matrix
    branch_map = _make_branch_map(case)

    gen_injections = JuMP.@expression(m, gen_map * m[:pg])
    line_injections = JuMP.@expression(m, branch_map * m[:pf])
    injections = JuMP.@expression(m, gen_injections + line_injections)
    if load_shed_enabled
        injections = JuMP.@expression(m, injections + sets.load_bus_map * m[:load_shed])
    end
    withdrawals = JuMP.@expression(m, bus_demand)
    if storage.enabled
        storage_bus_idx = [sets.bus_id2idx[b] for b in storage.gen[:, 1]]
        storage_map = sparse(
            storage_bus_idx, sets.storage_idx, 1, sets.num_bus, sets.num_storage
        )::SparseMatrixCSC
        injections = JuMP.@expression(m, injections + storage_map * m[:storage_dis])
        withdrawals = JuMP.@expression(m, withdrawals + storage_map * m[:storage_chg])
    end
    if demand_flexibility.enabled
        injections = JuMP.@expression(
            m, injections + sets.flexible_load_bus_map * m[:load_shift_dn]
        )
        withdrawals = JuMP.@expression(
            m, withdrawals + sets.flexible_load_bus_map * m[:load_shift_up]
        )
    end
    JuMP.@constraint(m, powerbalance, (injections .== withdrawals))
    println("powerbalance, setting names: ", Dates.now())
    interval_length = size(bus_demand)[2]
    for i in sets.bus_idx, j in 1:interval_length
        JuMP.set_name(
            powerbalance[i, j], "powerbalance[" * string(i) * "," * string(j) * "]"
        )
    end
end

function _add_constraint_load_shed!(
    m::JuMP.Model,
    case::Case,
    sets::Sets,
    demand_flexibility::DemandFlexibility,
    bus_demand::Matrix,
)
    interval_length = size(bus_demand)[2]
    demand_for_load_shed = JuMP.@expression(
        m,
        [i = 1:(sets.num_load_bus), j = 1:interval_length],
        bus_demand[sets.load_bus_idx[i], j],
    )
    if demand_flexibility.enabled
        flexible_load_bus_idx = indexin(sets.flexible_bus_idx, sets.load_bus_idx)
        flexible_load_map = sparse(
            flexible_load_bus_idx,
            1:(sets.num_flexible_bus),
            1,
            sets.num_load_bus,
            sets.num_flexible_bus,
        )::SparseMatrixCSC
        demand_for_load_shed = JuMP.@expression(
            m,
            demand_for_load_shed + flexible_load_map * m[:load_shift_up] -
                flexible_load_map * m[:load_shift_dn],
        )
    end
    JuMP.@constraint(
        m,
        load_shed_ub[i in 1:(sets.num_load_bus), j in 1:interval_length],
        m[:load_shed][i, j] <= demand_for_load_shed[i, j],
    )
end

function _add_constraints_storage_operation!(
    m::JuMP.Model,
    case::Case,
    sets::Sets,
    storage::Storage,
    interval_length::Int,
    storage_e0::Array{Float64,1},
)
    num_hour = interval_length

    println("storage soc_tracking: ", Dates.now())
    JuMP.@constraint(
        m,
        soc_tracking[i in 1:(sets.num_storage), h in 1:(num_hour - 1)],
        m[:storage_soc][i, h + 1] == (
            m[:storage_soc][i, h] * (1 - storage.sd_table.LossFactor[i]) +
            storage.sd_table.InEff[i] * m[:storage_chg][i, h + 1] -
            (1 / storage.sd_table.OutEff[i]) * m[:storage_dis][i, h + 1]
        ),
        container = Array,
    )
    println("storage initial_soc: ", Dates.now())
    JuMP.@constraint(
        m,
        initial_soc[i in 1:(sets.num_storage)],
        m[:storage_soc][i, 1] == (
            storage_e0[i] + storage.sd_table.InEff[i] * m[:storage_chg][i, 1] -
            (1 / storage.sd_table.OutEff[i]) * m[:storage_dis][i, 1]
        ),
        container = Array,
    )
    println("storage final_soc_min: ", Dates.now())
    JuMP.@constraint(
        m,
        soc_terminal_min[i in 1:(sets.num_storage)],
        m[:storage_soc][i, num_hour] >= storage.sd_table.ExpectedTerminalStorageMin[i],
        container = Array,
    )
    println("storage final_soc_max: ", Dates.now())
    JuMP.@constraint(
        m,
        soc_terminal_max[i in 1:(sets.num_storage)],
        m[:storage_soc][i, num_hour] <= storage.sd_table.ExpectedTerminalStorageMax[i],
        container = Array,
    )
end

function _add_constraints_demand_flexibility!(
    m::JuMP.Model,
    case::Case,
    sets::Sets,
    demand_flexibility::DemandFlexibility,
    interval_length::Int,
    init_shifted_demand::Array{Float64,1}=Float64[],
)
    if demand_flexibility.rolling_balance
        println("rolling load balance, first window: ", Dates.now())
        JuMP.@constraint(
            m,
            rolling_load_balance_first[i in 1:(sets.num_load_bus)],
            sum(
                m[:load_shift_up][i, j] - m[:load_shift_dn][i, j] for
                j in 1:(demand_flexibility.duration - 1)
            ) >= -1 * init_shifted_demand[i],
        )
        println("rolling load balance: ", Dates.now())
        JuMP.@constraint(
            m,
            rolling_load_balance[
                i in 1:(sets.num_flexible_bus),
                k in 1:(interval_length - demand_flexibility.duration + 1),
            ],
            sum(
                m[:load_shift_up][i, j] - m[:load_shift_dn][i, j] for
                j in k:(k + demand_flexibility.duration - 1)
            ) >= 0,
        )
    end
    if demand_flexibility.interval_balance
        println("interval load balance: ", Dates.now())
        JuMP.@constraint(
            m,
            interval_load_balance[i in 1:(sets.num_flexible_bus)],
            sum(
                m[:load_shift_up][i, j] - m[:load_shift_dn][i, j] for j in 1:interval_length
            ) >= -1 * init_shifted_demand[i],
        )
    end
end

function _add_constraints_initial_ramping!(
    m::JuMP.Model, case::Case, sets::Sets, initial_ramp_g0
)
    noninf_ramp_idx = findall(case.gen_ramp30 .!= Inf)
    println("initial rampup: ", Dates.now())
    JuMP.@constraint(
        m,
        initial_rampup[i in noninf_ramp_idx],
        m[:pg][i, 1] - initial_ramp_g0[i] <= case.gen_ramp30[i] * 2,
    )
    println("initial rampdown: ", Dates.now())
    JuMP.@constraint(
        m,
        initial_rampdown[i in noninf_ramp_idx],
        case.gen_ramp30[i] * -2 <= m[:pg][i, 1] - initial_ramp_g0[i],
    )
end

function _add_constraints_ramping!(
    m::JuMP.Model, case::Case, sets::Sets, interval_length::Int
)
    noninf_ramp_idx = findall(case.gen_ramp30 .!= Inf)
    println("rampup: ", Dates.now())
    JuMP.@constraint(
        m,
        rampup[i in noninf_ramp_idx, h in 1:(interval_length - 1)],
        m[:pg][i, h + 1] - m[:pg][i, h] <= case.gen_ramp30[i] * 2,
    )
    println("rampdown: ", Dates.now())
    JuMP.@constraint(
        m,
        rampdown[i in noninf_ramp_idx, h in 1:(interval_length - 1)],
        case.gen_ramp30[i] * -2 <= m[:pg][i, h + 1] - m[:pg][i, h],
    )
end

function _add_constraints_generator_segments!(
    m::JuMP.Model, case::Case, sets::Sets, hour_idx
)
    segment_width = (case.gen_pmax - case.gen_pmin) ./ sets.num_segments
    println("segment_max: ", Dates.now())
    JuMP.@constraint(
        m,
        segment_max[i in sets.noninf_pmax, s in sets.segment_idx, h in hour_idx],
        m[:pg_seg][i, s, h] <= segment_width[i],
    )
    println("segment_add: ", Dates.now())
    JuMP.@constraint(
        m,
        segment_add[i in sets.noninf_pmax, h in hour_idx],
        m[:pg][i, h] == case.gen_pmin[i] + sum(m[:pg_seg][i, :, h]),
    )
end

function _add_constraints_branch_flow_limits!(
    m::JuMP.Model, case::Case, sets::Sets, trans_viol_enabled, hour_idx
)
    branch_pmin = vcat(-1 * case.branch_rating, case.dcline_pmin)
    branch_pmax = vcat(case.branch_rating, case.dcline_pmax)
    if trans_viol_enabled
        JuMP.@expression(m, branch_limit_pmin, branch_pmin .- m[:trans_viol])
        JuMP.@expression(m, branch_limit_pmax, branch_pmax .+ m[:trans_viol])
    else
        JuMP.@expression(m, branch_limit_pmin, repeat(branch_pmin, 1, length(hour_idx)))
        JuMP.@expression(m, branch_limit_pmax, repeat(branch_pmax, 1, length(hour_idx)))
    end
    println("branch_min, branch_max: ", Dates.now())
    JuMP.@constraint(
        m,
        branch_min[br in sets.noninf_branch_idx, h in hour_idx],
        branch_limit_pmin[br, h] <= m[:pf][br, h],
    )
    println("branch_max: ", Dates.now())
    JuMP.@constraint(
        m,
        branch_max[br in sets.noninf_branch_idx, h in hour_idx],
        m[:pf][br, h] <= branch_limit_pmax[br, h],
    )
end

function _add_branch_angle_constraints!(m::JuMP.Model, case::Case, sets::Sets, hour_idx)
    # Explicit numbering here so that we constrain AC branches but not DC
    JuMP.@constraint(
        m,
        branch_angle[br in 1:(sets.num_branch_ac), h in hour_idx],
        (
            case.branch_reactance[br] * m[:pf][br, h] ==
            (m[:theta][sets.branch_to_idx[br], h] - m[:theta][sets.branch_from_idx[br], h])
        )
    )
end

function _add_profile_generator_limits!(
    m::JuMP.Model, case::Case, sets::Sets, hour_idx, start_index, interval_length
)
    end_index = start_index + interval_length - 1
    # Generation segments
    simulation_hydro = Matrix(case.hydro[start_index:end_index, 2:end])
    simulation_solar = Matrix(case.solar[start_index:end_index, 2:end])
    simulation_wind = Matrix(case.wind[start_index:end_index, 2:end])
    println("hydro_fixed: ", Dates.now())
    JuMP.@constraint(
        m,
        hydro_fixed[i in 1:(sets.num_hydro), h in hour_idx],
        m[:pg][sets.gen_hydro_idx[i], h] == simulation_hydro[h, i],
    )
    println("solar_max: ", Dates.now())
    JuMP.@constraint(
        m,
        solar_max[i in 1:(sets.num_solar), h in hour_idx],
        m[:pg][sets.gen_solar_idx[i], h] <= simulation_solar[h, i],
    )
    println("wind_max: ", Dates.now())
    JuMP.@constraint(
        m,
        wind_max[i in 1:(sets.num_wind), h in hour_idx],
        m[:pg][sets.gen_wind_idx[i], h] <= simulation_wind[h, i],
    )
end

function _add_objective_function!(
    m::JuMP.Model,
    case::Case,
    sets::Sets,
    storage::Storage,
    start_index::Int,
    end_index::Int,
    interval_length::Int,
    load_shed_enabled::Bool,
    load_shed_penalty::Number,
    trans_viol_enabled::Bool,
    trans_viol_penalty::Number,
    storage_e0::Array{Float64,1},
    demand_flexibility::DemandFlexibility,
)
    # Positional indices from mpc.gencost
    COST = 5
    fixed_cost = case.gencost[:, COST + 1]
    segment_width = (case.gen_pmax - case.gen_pmin) ./ sets.num_segments
    segment_slope = _build_segment_slope(case, sets.segment_idx, segment_width)

    # Start with generator variable O & M, piecewise
    obj = JuMP.@expression(
        m, sum(segment_slope[sets.noninf_pmax, :] .* m[:pg_seg][sets.noninf_pmax, :, :]),
    )
    # Add fixed costs
    JuMP.add_to_expression!(obj, JuMP.@expression(m, interval_length * sum(fixed_cost)))
    # Add load shed penalty (if necessary)
    if load_shed_enabled
        JuMP.add_to_expression!(
            obj, JuMP.@expression(m, load_shed_penalty * sum(m[:load_shed]))
        )
    end
    # Add transmission violation penalty (if necessary)
    if trans_viol_enabled
        JuMP.add_to_expression!(
            obj, JuMP.@expression(m, trans_viol_penalty * sum(m[:trans_viol]))
        )
    end
    # Pay for ending with less storage energy than initial
    if storage.enabled
        storage_penalty = JuMP.@expression(
            m,
            sum(
                (storage_e0 - m[:storage_soc][:, end]) .*
                storage.sd_table.TerminalStoragePrice,
            ),
        )
        JuMP.add_to_expression!(obj, storage_penalty)
    end

    # Pay for the cost of DR programs based on committed flexible demand
    if demand_flexibility.enabled
        # cost for increasing flexible load
        if !isnothing(demand_flexibility.cost_up)
            bus_demand_flex_cost_up = permutedims(
                Matrix(demand_flexibility.cost_up[start_index:end_index, 2:end])
            )
            demand_response_penalty_up = JuMP.@expression(
                m, sum(sum(bus_demand_flex_cost_up .* m[:load_shift_up]))
            )
            JuMP.add_to_expression!(obj, demand_response_penalty_up)
        end
        # cost for decreasing flexible load
        if !isnothing(demand_flexibility.cost_dn)
            bus_demand_flex_cost_dn = permutedims(
                Matrix(demand_flexibility.cost_dn[start_index:end_index, 2:end])
            )
            demand_response_penalty_dn = JuMP.@expression(
                m, sum(sum(bus_demand_flex_cost_dn .* m[:load_shift_dn]))
            )
            JuMP.add_to_expression!(obj, demand_response_penalty_dn)
        end
    end
    # Finally, set as objective of model
    JuMP.@objective(m, Min, obj)
end

"""
    _build_model(m; case=case, storage=storage, start_index=x,
                 interval_length=y[, kwargs...])

Given a Case object and a set of options, build an optimization model.
Returns a JuMP.Model instance.
"""
function _build_model(
    m::JuMP.Model;
    case::Case,
    storage::Storage,
    demand_flexibility::DemandFlexibility,
    start_index::Int,
    interval_length::Int,
    demand_scaling::Number=1.0,
    load_shed_enabled::Bool=false,
    load_shed_penalty::Number=9000,
    trans_viol_enabled::Bool=false,
    trans_viol_penalty::Number=100,
    initial_ramp_enabled::Bool=false,
    initial_ramp_g0::Array{Float64,1}=Float64[],
    storage_e0::Array{Float64,1}=Float64[],
    init_shifted_demand::Array{Float64,1}=Float64[],
)::JuMP.Model
    # Positional indices from mpc.gen
    PMAX = 9
    PMIN = 10
    println("building sets: ", Dates.now())
    # Sets - time periods
    hour_idx = 1:interval_length
    end_index = start_index + interval_length - 1
    # Sets - static
    sets = _make_sets(case; storage=storage, demand_flexibility=demand_flexibility)
    println("parameters: ", Dates.now())
    # Parameters
    bus_demand = _make_bus_demand(case, start_index, end_index) * demand_scaling
    # Demand flexibility parameters (if present)
    if demand_flexibility.enabled
        (bus_demand_flex_amt_up, bus_demand_flex_amt_dn) = _make_bus_demand_flexibility_amount(
            case, demand_flexibility, start_index, end_index
        )
    end

    println("variables: ", Dates.now())
    # Variables
    # Explicitly declare containers as Array of VariableRefs, not DenseArrayAxis
    JuMP.@variables(
        m,
        begin
            pg[sets.gen_idx, hour_idx] >= 0, (container = Array)
            pg_seg[sets.gen_idx, sets.segment_idx, hour_idx] >= 0, (container = Array)
            pf[sets.branch_idx, hour_idx], (container = Array)
            theta[sets.bus_idx, hour_idx], (container = Array)
        end
    )
    if load_shed_enabled
        JuMP.@variable(
            m,
            load_shed[i in 1:(sets.num_load_bus), j in 1:interval_length] >= 0,
            container = Array
        )
    end
    if trans_viol_enabled
        JuMP.@variable(
            m, 0 <= trans_viol[i in sets.branch_idx, j in hour_idx], container = Array
        )
    end
    if storage.enabled
        storage_max_dis = storage.gen[:, PMAX]
        storage_max_chg = -1 * storage.gen[:, PMIN]
        storage_min_energy = storage.sd_table.MinStorageLevel
        storage_max_energy = storage.sd_table.MaxStorageLevel
        JuMP.@variables(
            m,
            begin
                (
                    0 <=
                    storage_chg[i in sets.storage_idx, j in hour_idx] <=
                    storage_max_chg[i]
                ),
                (container = Array)
                (
                    0 <=
                    storage_dis[i in sets.storage_idx, j in hour_idx] <=
                    storage_max_dis[i]
                ),
                (container = Array)
                (
                    storage_min_energy[i] <=
                    storage_soc[i in sets.storage_idx, j in hour_idx] <=
                    storage_max_energy[i]
                ),
                (container = Array)
            end
        )
    end
    if demand_flexibility.enabled
        # The amount of demand that is curtailed from the base load
        JuMP.@variable(
            m,
            0 <=
                load_shift_dn[i in 1:(sets.num_flexible_bus), j in 1:interval_length] <=
                bus_demand_flex_amt_dn[i, j],
        )
        # The amount of demand that is added from the base load
        JuMP.@variable(
            m,
            0 <=
                load_shift_up[i in 1:(sets.num_flexible_bus), j in 1:interval_length] <=
                bus_demand_flex_amt_up[i, j],
        )
    end

    println("constraints: ", Dates.now())
    # Constraints

    println("powerbalance: ", Dates.now())
    _add_constraint_power_balance!(
        m, case, sets, storage, demand_flexibility, bus_demand, load_shed_enabled
    )

    if load_shed_enabled
        _add_constraint_load_shed!(m, case, sets, demand_flexibility, bus_demand)
    end

    if storage.enabled
        _add_constraints_storage_operation!(
            m, case, sets, storage, interval_length, storage_e0
        )
    end

    if demand_flexibility.enabled
        _add_constraints_demand_flexibility!(
            m, case, sets, demand_flexibility, interval_length, init_shifted_demand
        )
    end

    if initial_ramp_enabled
        _add_constraints_initial_ramping!(m, case, sets, initial_ramp_g0)
    end
    if length(hour_idx) > 1
        _add_constraints_ramping!(m, case, sets, interval_length)
    end

    _add_constraints_generator_segments!(m, case, sets, hour_idx)

    _add_constraints_branch_flow_limits!(m, case, sets, trans_viol_enabled, hour_idx)

    println("branch_angle: ", Dates.now())
    _add_branch_angle_constraints!(m, case, sets, hour_idx)

    # Constrain variable generators based on profiles
    _add_profile_generator_limits!(m, case, sets, hour_idx, start_index, interval_length)

    println("objective: ", Dates.now())
    _add_objective_function!(
        m,
        case,
        sets,
        storage,
        start_index,
        end_index,
        interval_length,
        load_shed_enabled,
        load_shed_penalty,
        trans_viol_enabled,
        trans_viol_penalty,
        storage_e0,
        demand_flexibility,
    )

    println(Dates.now())
    return m
end


text:
module SSLTestWorldclim
using SimpleSDMLayers
using Test

wc1and2 = worldclim([1,2])
@test typeof(first(wc1and2)) <: SimpleSDMPredictor

temp = worldclim(1)
@test size(temp) == (1080, 2160)
@test round(first(stride(temp)); digits=2) ≈ round(last(stride(temp)); digits=2)
@test length(longitudes(temp)) == 2160
@test length(latitudes(temp)) == 1080

wc3 = worldclim(3)
@test typeof(wc3) <: SimpleSDMPredictor

wcrange = worldclim(1:5)
@test eltype(wcrange) <: SimpleSDMPredictor

end


text:
using OdinSon
using AffineTransforms
using Distributions
import Base: *, \
# extend the algebra of AffineTransform to be more permissive, tform{fwd, inv}
# does type checking, sent a fix upstream, remove once new tagged version
*(a::AffineTransform, v::AbstractArray) = tformfwd(a, v)
\(a::AffineTransform, x::AbstractArray) = tforminv(a, x)
#TODO: implement a real function like Mathematica's ReflectionMatrix
tformreflect = AffineTransform([-1.0 0.0; 0.0 1.0], [0.0, 0.0])

# wow , \theta vs 0 is very hard to see in this font
const rots = map(Θ -> tformrotate(Θ), 0:π/3:(2π - π/3))

function snowflake(pt)
    #TODO: this will be slow do to column major order, which I could easily switch
    mapslices(x -> rots[1]*x, pts, 2)

    # to decode: Join[Map[ReflectionMatrix[{1, 0}].# &, #], #] &
    endline = vcat([0.0 0.0], pts)
    startline = mapslices(x->tformreflect*x, endline, 2)
    arg2 = vcat(startline, endline)

    # this recreates the Outer(#1 . #2&, rots, arg2, 1) call in mathematica
    ##TODO: out can contain an 2dim Array, so the size is not enforced. I would likely
    ## need to use FixedSizeArrays to do this perfectly
    out = Array{Array{Float64, 2}}(length(rots))
    for i = 1:length(rots)
        poly = zeros(2, size(arg2, 1))
        for j = 1:size(arg2, 1)
            poly[:, j] = rots[i] * arg2[j:j, :]'
        end
        out[i] = poly
    end
    ps = map(p->Polygon(p, style=Style(stroke=nothing, fill=NC"white", fill_opacity=0.5)), out)

    return Canvas(ps, style=Style(fill=NC"black"))
end

pts = rand(Uniform(-1, 1), (rand(3:9), 2))
flake = snowflake(pts)
OdinSon.render(flake)

max_stars_repo_name: gabrielgellner/OdinSon

max_issues_repo_name: gabrielgellner/OdinSon

max_forks_repo_name: gabrielgellner/OdinSon


text:
@inline log(x) = Base.log(x)

@inline function log(x::T) where T <:Real
        _0 = zero(x)
        ifelse(x>=zero(x),Base.log(max(_0,x)),_0/_0)
    end

@inline sqrt(x) = Base.sqrt(x)
@inline function sqrt(x::T) where T <:Real
    _0 = zero(x)
    ifelse(x>=zero(x),Base.sqrt(max(_0,x)),_0/_0)
end


text:
"""
    Parameters(N::Int64, M::Int64, u::Array{Float64,2}, Rm::Vector{Float64}, l::Array{Float64,2}, l_sum::Vector{Float64}, ρ::Vector{Float64})

Type containing the parameters for a simuation.

# Arguments

- `N::Int64` : Number of consumers
- `M::Int64`: Number of Resourcces
- `u::Array{Float64,2}`: Array of uptake rates of resource i by consumer j
- `Rm::Array{Float64,1}`: Vector of maintennce respiration costs for consumers
- `l::Array{Float64,2}`: Array of the proportion of uptake of resource i released as resource j
- `l_sum::Array{Float64,1}`: Array of the total proportion of uptake of resource j released
- `ρ::Array{Float64,1}`: Vector of inflow rates for resources
- `ω::Array{Float64,1}`: Vector of outflow rates for resources
"""
mutable struct Parameters
    N::Int64
    M::Int64
    u::Array{Float64,2}
    Rm::Array{Float64,1}
    l::Array{Float64,2}
    l_sum::Array{Float64,1}
    ρ::Array{Float64,1}
    ω::Array{Float64,1}
end

"""
    make_parameters(N::Int64, M::Int64, u::Array{Float64,2}, Rm::Vector{Float64}, l::Array{Float64,2}, ρ::Vector{Float64}, ω::Array{Float64,1})

Helper function that takes values for parameters and returns a `Parameters`
object. Also does checks internally to make sure the values are correct.
"""
function make_parameters(N::Int64, M::Int64,
        u::Array{Float64,2}, Rm::Vector{Float64},
        l::Array{Float64,2}, ρ::Vector{Float64},ω::Array{Float64,1})

    #assert parameters are the right size
    @assert all([N,M] .> 0) "Number of consumers and resources must both be greater than 0"
    @assert size(u) == (N,M) "Uptake array is the wrong size"
    @assert length(Rm) == N "Rm vector is the wrong length"
    @assert size(l) == (M,M) "Leakage array is the wrong size"
    @assert length(ρ) == M "Inflow vector is the wrong length"

    #assert that the leakage proportions add up to 1
    @assert all(sum(l,dims = 2) .<= 1) "One or more l values are > 1"

    return( Parameters(copy(N), copy(M), copy(u), copy(Rm), copy(l), sum(copy(l), dims=2)[:], copy(ρ), copy(ω)))
end

"""
make_parameters(N::Int64, M::Int64, u::Array{Float64,2}, Rm::Vector{Float64}, l::Array{Float64,2}, ρ::Vector{Float64})

creates parameters with ω = 0 for all resources
"""
function make_parameters(N::Int64, M::Int64,
        u::Array{Float64,2}, Rm::Vector{Float64},
        l::Array{Float64,2}, ρ::Vector{Float64})

    #assert parameters are the right size
    @assert all([N,M] .> 0) "Number of consumers and resources must both be greater than 0"
    @assert size(u) == (N,M) "Uptake array is the wrong size"
    @assert length(Rm) == N "Rm vector is the wrong length"
    @assert size(l) == (M,M) "Leakage array is the wrong size"
    @assert length(ρ) == M "Inflow vector is the wrong length"

    #assert that the leakage proportions add up to 1
    @assert all(sum(l,dims = 2) .<= 1) "One or more l values are > 1"

    return( Parameters(copy(N), copy(M), copy(u), copy(Rm), copy(l), sum(copy(l), dims=2)[:], copy(ρ), zeros(M)))
end