text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

name = "MongoC"
version = v"1.16.2"

# Collection of sources required to complete build
sources = [
    ArchiveSource("https://github.com/mongodb/mongo-c-driver/releases/download/1.16.2/mongo-c-driver-1.16.2.tar.gz", "0a722180e5b5c86c415b9256d753b2d5552901dc5d95c9f022072c3cd336887e")
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir
cd mongo-c-driver-1.16.2/
sed -i "s/Windows.h/windows.h/" src/libmongoc/src/mongoc/mongoc-client.c
sed -i "s/WinDNS.h/windns.h/" src/libmongoc/src/mongoc/mongoc-client.c
sed -i "s/Mstcpip.h/mstcpip.h/" src/libmongoc/src/mongoc/mongoc-client.c
sed -i "s/Mstcpip.h/mstcpip.h/" src/libmongoc/src/mongoc/mongoc-socket.c
sed -i "s/Dnsapi/dnsapi/" build/cmake/FindResSearch.cmake
mkdir cmake-build
cd cmake-build/

if [[ "${nbits}" == 32 ]]; then
    export CFLAGS="-Wl,-rpath-link,/opt/${target}/${target}/lib"
elif [[ "${target}" != *-apple-* ]]; then
    export CFLAGS="-Wl,-rpath-link,/opt/${target}/${target}/lib64"
fi

cmake  -DENABLE_AUTOMATIC_INIT_AND_CLEANUP=OFF -DCMAKE_BUILD_TYPE=Release -DENABLE_SSL=OPENSSL -DENABLE_SASL=OFF -DENABLE_EXAMPLES=OFF -DENABLE_TESTS=OFF -DENABLE_UNINSTALL=OFF -DENABLE_STATIC=OFF -DENABLE_SNAPPY=ON -DENABLE_MONGOC=ON -DENABLE_BSON=ON -DOPENSSL_ROOT_DIR=$prefix -DCMAKE_INSTALL_PREFIX=$prefix -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TARGET_TOOLCHAIN} ..
make -j${nproc}
make install

if [[ "${target}" == *-apple-darwin* ]]; then
    rm ${libdir}/libbson-1.0.0.dylib
    rm ${libdir}/libmongoc-1.0.0.dylib
    mv ${libdir}/libmongoc-1.0.0.0.0.dylib ${libdir}/libmongoc-1.0.0.dylib
    mv ${libdir}/libbson-1.0.0.0.0.dylib ${libdir}/libbson-1.0.0.dylib
fi
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = supported_platforms()

# The products that we will ensure are always built
products = [
    LibraryProduct(["libbson-1", "libbson-1.0", "libbson"], :libbson),
    LibraryProduct(["libmongoc-1", "libmongoc-1.0", "libmongoc"], :libmongoc),
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency(PackageSpec(name="OpenSSL_jll", uuid="458c3c95-2e84-50aa-8efc-19380b2a3a95"))
    Dependency(PackageSpec(name="Zlib_jll", uuid="83775a58-1f1d-513f-b197-d71354ab007a"))
    Dependency(PackageSpec(name="Zstd_jll", uuid="3161d3a3-bdf6-5164-811a-617609db77b4"))
    Dependency(PackageSpec(name="snappy_jll", uuid="fe1e1685-f7be-5f59-ac9f-4ca204017dfd"))
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)


text:
@testset "radial basis function" begin
    xc = rand(10)
    yc = rand(10)
    e = rand(10)
    c = rand(10)
    d = rand(3)
    r = RBF2D(xc, yc; c=c, eps=e, d=d, kind = 0)

    x = rand(5)
    y = rand(5)
    o = r(x, y)

    O = zeros(5)
    for i = 1:5
        for j = 1:10
            dist = sqrt((x[i]-xc[j])^2 + (y[i]-yc[j])^2)
            O[i] += c[j] * exp(-(e[j]*dist)^2)
        end
        O[i] += d[1] + d[2] * x[i] + d[3] * y[i]
    end
    init(sess)
    @test norm(run(sess, o)-O)<1e-5
end

@testset "interp1" begin
    Random.seed!(233)
    x = sort(rand(10))
    y = @. x^2 + 1.0
    z = [x[1]; x[2]; rand(5) * (x[end]-x[1]) .+ x[1]; x[end]]
    u = interp1(x,y,z)
    @test norm(run(sess, u)-[1.026422850882909
                    1.044414684090653
                    1.312604319732756
                    1.810845361128137
                    1.280789421523103
                    1.600084940795178
                    1.930560200260898
                    1.972130181835701]) < 1e-10
end

@testset "rbf3d" begin

    function f0(r, e)
        return exp(-(e*r)^2)
    end

    function f1(r, e)
        return sqrt((1+(e*r)^2))
    end

    function f2(r, e)
        return 1/(1+(e*r)^2)
    end

    function f3(r, e)
        return 1/sqrt(1+(e*r)^2)
    end

    fs = [f0, f1, f2, f3]
    nc = 100
    n = 10
    xc = rand(nc)
    yc = rand(nc)
    zc = rand(nc)
    x = rand(n)
    y = rand(n)
    z = rand(n)
    c = rand(nc)
    d = rand(4)
    e = rand(nc)

    for kind = 1:4
        out = zeros(n)
        for i = 1:n
            r = @. sqrt((x[i] - xc)^2 + (y[i] - yc)^2 + (z[i] - zc)^2)
            out[i] = sum(c .* fs[kind].(r, e)) + d[1] + d[2] * x[i] + d[3] * y[i] + d[4] * z[i]
        end

        rbf = RBF3D(xc, yc, zc, c = c, eps = e, d = d, kind = kind - 1)
        o = rbf(x,y,z)

        o_ = run(sess, o)

        @test maximum(abs.(out .- o_)) < 1e-8
    end

    for kind = 1:4
        out = zeros(n)
        for i = 1:n
            r = @. sqrt((x[i] - xc)^2 + (y[i] - yc)^2 + (z[i] - zc)^2)
            out[i] = sum(c .* fs[kind].(r, e)) + d[1]
        end

        rbf = RBF3D(xc, yc, zc, c = c, eps = e, d = d[1:1], kind = kind - 1)
        o = rbf(x,y,z)

        o_ = run(sess, o)

        @test maximum(abs.(out .- o_)) < 1e-8
    end

    for kind = 1:4
        out = zeros(n)
        for i = 1:n
            r = @. sqrt((x[i] - xc)^2 + (y[i] - yc)^2 + (z[i] - zc)^2)
            out[i] = sum(c .* fs[kind].(r, e))
        end

        rbf = RBF3D(xc, yc, zc, c = c, eps = e, kind = kind - 1)
        o = rbf(x,y,z)

        o_ = run(sess, o)

        @test maximum(abs.(out .- o_)) < 1e-8
    end
end

ext: jl


text:
using AHPQ
using Test

@testset "testmetrics" begin
    # Data Generation
    a = Array{Int}(undef, 10, 4)
    for i in 1:4 a[:,i] = collect(1:10) end
    b = deepcopy(a)

    @info("Testing RecallFull")
    @test recallfull(a, b)==1.0
    b[9,1] = 12

    recallfull(a,b)
    @test recallfull(a, b)==0.975

    @info("Testing RecallN")
    @test recallN(a, b, 10) == 0.975
    @test recallN(a, b, 5) == 1.0

    @info("Testing Recall1atN")
    c = Array{Int}(undef, 10, 4)
    for i in 1:4 c[:,i] = collect(10:19) end
    c[5,1] = 1
    @test recall1atN(c, a, 10) == 0.25
    @test recall1atN(c, a, 4) == 0
    @test recall1atN(b, a, 1) == 1

    @test get1atNscores(c, a, 5) == [0.0, 0.0, 0.0, 0.0, 0.25]

    @info("Testing Approx_error")

    # Data prep
    n_dp = 100
    n_dim = 4
    traindata = rand(n_dim, n_dp)
    queries = rand(n_dim, 10)
    innerproducts = traindata' * queries
    n_neighbors=10
    rtrue = mapslices(x -> partialsortperm(x, 1:n_neighbors, rev=true), innerproducts,dims=1);

    # Train AHPQ
    ahpq = builder(deepcopy(traindata); T=0, n_codebooks=1, n_centers=2, training_points=0)
    approx_error(ahpq.qd, deepcopy(traindata), rtrue, queries)

    @test true
end


text:
"""
    struct NoPlasticity end
"""
struct NoPlasticity end
export NoPlasticity
updateweights!(net, i) = Nothing

"""
    mutable struct TeacherSignalPlasticity
        islearning::Bool
        isweightplastic::Array{Bool, 1}
        plasticweights::Array{Int64, 1}
        targets::Array{Float64, 1}
        learningrate::Float64
        weightdecaybound::Float64
    end
"""
@def plasticity_common_fields begin
    islearning::Bool
    isweightplastic::Array{Bool, 1}
    plasticweights::Array{Int64, 1}
    targets::Array{Float64, 1}
    learningrate::Float64
    weightdecaybound::Float64
end
mutable struct TeacherSignalPlasticity
    @plasticity_common_fields
end
export TeacherSignalPlasticity
function TeacherSignalPlasticity(isweightplastic, targets, learningrate; weightdecaybound = Inf)
    TeacherSignalPlasticity(true, isweightplastic, findall(isweightplastic),
                            targets, learningrate, weightdecaybound)
end

"""
    mutable struct TeacherSignalPlasticitySTDP
        islearning::Bool
        isweightplastic::Array{Bool, 1}
        plasticweights::Array{Int64, 1}
        targets::Array{Float64, 1}
        learningrate::Float64
        weightdecaybound::Float64
        lastspiketimesatsynapses::Array{Float64, 1}
        pretracesatlastupdate::Array{Float64, 1}
        posttracesatlastupdate::Array{Float64, 1}
    end
"""
mutable struct TeacherSignalPlasticitySTDP
    @plasticity_common_fields
    lastspiketimesatsynapses::Array{Float64, 1} #time of last spike happening at every synapse
    pretracesatlastupdate::Array{Float64, 1} #trace (as pre-neuron) at last update (pre or post spike) at every synapse
    posttracesatlastupdate::Array{Float64, 1} #trace (as post-neuron) ''
end
export TeacherSignalPlasticitySTDP
function TeacherSignalPlasticitySTDP(isweightplastic, targets, learningrate; weightdecaybound = Inf)
    TeacherSignalPlasticitySTDP(true, isweightplastic, findall(isweightplastic),
                            targets, learningrate, weightdecaybound, zeros(length(isweightplastic)),
                            zeros(length(isweightplastic)), zeros(length(isweightplastic)))
end

################################################################################
# Plasticity Rules
################################################################################

@inline deltarule(lr, pre, post, target) = lr * (target - post) #* pre
export deltarule
@inline function STDPrule(lr, pretracesatlastupdate, posttracesatlastupdate, lastspiketimeatsynapse,
                          target, spiketracetimeconstant, nettime)
    lr * spiketracetimeconstant * pretracesatlastupdate *
    (target * (1 - exp((lastspiketimeatsynapse - nettime) / spiketracetimeconstant))
    - posttracesatlastupdate * sqrt(pi)/2 * erf((nettime - lastspiketimeatsynapse) / spiketracetimeconstant))
    #- posttracesatlastupdate * sqrt(pi) * erf((nettime - lastspiketimeatsynapse) / spiketracetimeconstant))
end

@inline function weightdecay(lr, pre, post, target, weightbound)
  (1 - lr / weightbound)
end

################################################################################
# Weight update
################################################################################

@inline function updateweights!(net::AbstractNet{TeacherSignalPlasticity}, i)
    !net.plasticityrule.islearning && return 0.
    for post in net.posts[i]
        if net.plasticityrule.isweightplastic[post[2]]
            net.weights[post[2]] += deltarule(net.plasticityrule.learningrate,
                                              net.spiketrace[i],
                                              net.spiketrace[post[1]],
                                              net.plasticityrule.targets[post[1]])
            (net.plasticityrule.weightdecaybound != Inf) &&
            (net.weights[post[2]] *= weightdecay(net.plasticityrule.learningrate, net.spiketrace[i], net.spiketrace[post[1]],
                                            net.plasticityrule.targets[post[1]], net.plasticityrule.weightdecaybound))
        end
    end
end

@inline function updateweights!(net::AbstractNet{TeacherSignalPlasticitySTDP}, i)
    for post in net.posts[i] #update all synapses to postsynaptic neurons
        if net.plasticityrule.isweightplastic[post[2]]
            if net.plasticityrule.islearning
              net.weights[post[2]] += STDPrule(net.plasticityrule.learningrate,
                                                net.plasticityrule.pretracesatlastupdate[post[2]], #pretrace at last spike
                                                net.plasticityrule.posttracesatlastupdate[post[2]], #post ''
                                                net.plasticityrule.lastspiketimesatsynapses[post[2]], #time of last spike
                                                net.plasticityrule.targets[post[1]],
                                                net.parameters.spiketrace_timeconstant,
                                                net.t)
            end
            net.plasticityrule.pretracesatlastupdate[post[2]] = net.spiketrace[i] #Save present states and times into memory variables
            net.plasticityrule.posttracesatlastupdate[post[2]] = net.spiketrace[post[1]]
            net.plasticityrule.lastspiketimesatsynapses[post[2]] = net.t
        end
    end
    for pre in net.pres[i] #update all synapses to presynaptic neurons
        if net.plasticityrule.isweightplastic[pre[2]]
            if net.plasticityrule.islearning
              net.weights[pre[2]] += STDPrule(net.plasticityrule.learningrate,
                                                net.plasticityrule.pretracesatlastupdate[pre[2]],
                                                net.plasticityrule.posttracesatlastupdate[pre[2]],
                                                net.plasticityrule.lastspiketimesatsynapses[pre[2]],
                                                net.plasticityrule.targets[i],
                                                net.parameters.spiketrace_timeconstant,
                                                net.t)
            end
            net.plasticityrule.pretracesatlastupdate[pre[2]] = net.spiketrace[pre[1]]
            net.plasticityrule.posttracesatlastupdate[pre[2]] = net.spiketrace[i]
            net.plasticityrule.lastspiketimesatsynapses[pre[2]] = net.t
        end
    end
end


text:
# Activate test environment on older Julia versions
if VERSION < v"1.2"
    using Pkg: Pkg
    Pkg.activate(@__DIR__)
    Pkg.develop(Pkg.PackageSpec(; path=dirname(@__DIR__)))
    Pkg.instantiate()
end

using AbstractPPL
using Documenter
using Test

@testset "AbstractPPL.jl" begin
    include("deprecations.jl")

    @testset "doctests" begin
        DocMeta.setdocmeta!(
            AbstractPPL,
            :DocTestSetup,
            :(using AbstractPPL);
            recursive=true,
        )
        doctest(AbstractPPL; manual=false)
    end
end


text:
using JuMP
using SCS
using LinearAlgebra
using FFTW

N = 10
c = ones(N)
A = [
    zeros(1,N-2) 0;
    Matrix(I,N-2,N-2) zeros(N-2,1)
]
B = [1; zeros(N-2,1)]

model = Model(SCS.Optimizer)
set_silent(model)

@variable(model, S[1:N, 1:N], PSD)
@variable(model, P[1:N-1, 1:N-1], Symmetric)
@variable(model, r[1:N] >= 0)
# @variable(model, R_I[1:N])
@variable(model, R_R[1:N])
@variable(model, C[1:N-1])
@variable(model, D)
@variable(model, t >= 0)

# helper constraints
@constraint(model, [i = 0:N-2], r[i + 2] == r[N-i])  # r must be Symmetric

# oficial constraints
@constraint(model, r[2:N] .== C[1:N-1])
@constraint(model, r[1] == D)

@constraint(model, S[1:N-1, 1:N-1] .== P - transpose(A)*P*A)
@constraint(model, S[1:N-1, N] .== transpose(C) - transpose(A)*P*B)
@constraint(model, S[N, 1:N-1] .== C - transpose(B)*P*A)
@constraint(model, S[N,N] - 2*D + (transpose(B)*P*B)[1,1] == 0)

@constraint(model, [k = 1:N], sum(r[n]*cos(2*pi*(n-1)*(k-1)/N) for n in 1:N) == R_R[k])
# @constraint(model, [k = 1:N], -sum(r[n]*sin(2*pi*(n-1)*(k-1)/N) for n in 1:N) == R_I[k])

# @constraint(model, 1 - t <= R_R[1] <= 1 + t)
# @constraint(model, 1 - t <= R_R[2] <= 1 + t)
# @constraint(model, 1 - t <= R_R[N] <= 1 + t)

@constraint(model, R_R[1] + t >= 1)
@constraint(model, R_R[1] - t <= 1)
@constraint(model, R_R[2] + t >= 1)
@constraint(model, R_R[2] - t <= 1)
@constraint(model, R_R[N] + t >= 1)
@constraint(model, R_R[N] - t <= 1)

@constraint(model, R_R[5] + t >= 0)
@constraint(model, R_R[5] - t <= 0)
@constraint(model, R_R[6] + t >= 0)
@constraint(model, R_R[6] - t <= 0)
@constraint(model, R_R[7] + t >= 0)
@constraint(model, R_R[7] - t <= 0)

@objective(model, Min, t)
JuMP.optimize!(model)

print(model)

r = JuMP.value.(r)
R = fft(r)

obj_value = JuMP.objective_value(model)
println("Objective value: ", obj_value)
println("r: ")
show(stdout, MIME("text/plain"), r)
println()
println("R_R: ")
show(stdout, MIME("text/plain"), JuMP.value.(R_R))
# println()
# println("R_I: ")
# show(stdout, MIME("text/plain"), JuMP.value.(R_I))
println()
println()
println()

# println("A: ")
# show(stdout, MIME("text/plain"), A)
# println()
# println("B: ")
# show(stdout, MIME("text/plain"), B)
# println()
# println("C: ")
# show(stdout, MIME("text/plain"), JuMP.value.(C))
# println()
# println("D: ")
# show(stdout, MIME("text/plain"), JuMP.value(D))
# println()
# println()
# println("S: ")
# show(stdout, MIME("text/plain"), JuMP.value.(S))
# println()
# println("P: ")
# show(stdout, MIME("text/plain"), JuMP.value.(P))

max_stars_repo_name: troncosoae/opti-electric-project

max_issues_repo_name: troncosoae/opti-electric-project

max_forks_repo_name: troncosoae/opti-electric-project


text:
# Operator hierarchy
const AssignmentOp  = 1
const ConditionalOp = 2
const ArrowOp       = 3
const LazyOrOp      = 4
const LazyAndOp     = 5
const ComparisonOp  = 6
const PipeOp        = 7
const ColonOp       = 8
const PlusOp        = 9
const BitShiftOp    = 10
const TimesOp       = 11
const RationalOp    = 12
const PowerOp       = 13
const DeclarationOp = 14
const WhereOp       = 15
const DotOp         = 16
const PrimeOp       = 16
const DddotOp       = 7
const AnonFuncOp    = 14

@enum(Head,IDENTIFIER,
PUNCTUATION,
OPERATOR,
KEYWORD,
LITERAL,
NoHead,
Call,
UnaryOpCall,
BinaryOpCall,
WhereOpCall,
ConditionalOpCall,
ChainOpCall,
ColonOpCall,
Abstract,
Begin,
Block,
Braces,
BracesCat,
Const,
Comparison,
Curly,
Do,
Filter,
Flatten,
For,
FunctionDef,
Generator,
Global,
GlobalRefDoc,
If,
Kw,
Let,
Local,
Macro,
MacroCall,
MacroName,
Mutable,
Outer,
Parameters,
Primitive,
Quote,
Quotenode,
InvisBrackets,
StringH,
Struct,
Try,
TupleH,
FileH,
Return,
While,
x_Cmd,
x_Str,
ModuleH,
BareModule,
TopLevel,
Export,
Import,
ImportAll,
Using,
Comprehension,
DictComprehension,
TypedComprehension,
Hcat,
TypedHcat,
Ref,
Row,
Vcat,
TypedVcat,
Vect,
ErrorToken)

@enum(ErrorKind,
    UnexpectedToken,
    CannotJuxtapose,
    UnexpectedWhiteSpace,
    UnexpectedNewLine,
    ExpectedAssignment,
    UnexpectedAssignmentOp,
    MissingConditional,
    MissingCloser,
    InvalidIterator,
    StringInterpolationWithTrailingWhitespace,
    Unknown)

const NoKind = Tokenize.Tokens.begin_keywords

mutable struct Binding
    name::String
    val
    t
    refs::Vector
    overwrites::Union{Nothing,Binding}
end
function Binding(x)
    Binding(str_value(get_name(x)), x, nothing, [], nothing)
end

mutable struct Scope
    parent::Union{Nothing,Scope}
    names::Dict{String,Binding}
    modules::Union{Nothing,Dict{String,Any}}
    ismodule::Bool
end

Scope() = Scope(nothing, Dict{String,Binding}(), nothing, false)

mutable struct EXPR
    typ::Head
    args::Union{Nothing,Vector{EXPR}}
    fullspan::Int
    span::Int
    val::Union{Nothing,String}
    kind::Tokenize.Tokens.Kind
    dot::Bool
    parent::Union{Nothing,EXPR}
    scope::Union{Nothing,Scope}
    binding::Union{Nothing,Binding}
    ref
end

function EXPR(T::Head, args::Vector{EXPR}, fullspan::Int, span::Int)
    ex = EXPR(T, args, fullspan, span, nothing, NoKind, false, nothing, nothing, nothing, nothing)
    for c in args
        setparent!(c, ex)
    end
    ex
end

function EXPR(T::Head, args::Vector{EXPR})
    ret = EXPR(T, args, 0, 0)
    update_span!(ret)
    ret
end

@noinline mIDENTIFIER(ps::ParseState) = EXPR(IDENTIFIER, nothing, ps.nt.startbyte - ps.t.startbyte, ps.t.endbyte - ps.t.startbyte + 1, val(ps.t, ps), NoKind, false, nothing, nothing, nothing, nothing)

mPUNCTUATION(kind::Tokens.Kind, fullspan::Int, span::Int) = EXPR(PUNCTUATION, nothing, fullspan, span, nothing, kind, false, nothing, nothing, nothing, nothing)
@noinline mPUNCTUATION(ps::ParseState) = EXPR(PUNCTUATION, nothing, ps.nt.startbyte - ps.t.startbyte, ps.t.endbyte - ps.t.startbyte + 1, nothing, kindof(ps.t), false, nothing, nothing, nothing, nothing)

mOPERATOR(fullspan::Int, span::Int, kind::Tokens.Kind, dotop::Bool) = EXPR(OPERATOR, nothing, fullspan, span, nothing, kind, dotop, nothing, nothing, nothing, nothing)
@noinline mOPERATOR(ps::ParseState) = EXPR(OPERATOR, nothing, ps.nt.startbyte - ps.t.startbyte, ps.t.endbyte - ps.t.startbyte + 1, ps.t.suffix ? val(ps.t, ps) : nothing, kindof(ps.t), ps.t.dotop, nothing, nothing, nothing, nothing)

mKEYWORD(kind::Tokens.Kind, fullspan::Int, span::Int) = EXPR(KEYWORD, nothing, fullspan, span, nothing, kind, false, nothing, nothing, nothing, nothing)
@noinline mKEYWORD(ps::ParseState) = EXPR(KEYWORD, nothing, ps.nt.startbyte - ps.t.startbyte, ps.t.endbyte - ps.t.startbyte + 1, nothing, kindof(ps.t), false, nothing, nothing, nothing, nothing)

mLITERAL(fullspan::Int, span::Int, val::String, kind::Tokens.Kind) = EXPR(LITERAL, nothing, fullspan, span, val, kind, false, nothing, nothing, nothing, nothing)
@noinline function mLITERAL(ps::ParseState)
    if kindof(ps.t) == Tokens.STRING || kindof(ps.t) == Tokens.TRIPLE_STRING ||
        kindof(ps.t) == Tokens.CMD || kindof(ps.t) == Tokens.TRIPLE_CMD
        return parse_string_or_cmd(ps)
    else
        mLITERAL(ps.nt.startbyte - ps.t.startbyte, ps.t.endbyte - ps.t.startbyte + 1, val(ps.t, ps), kindof(ps.t))
    end
end

span(x::EXPR) = x.span

function update_span!(x::EXPR) end
function update_span!(x::EXPR)
    (x.args isa Nothing || isempty(x.args)) && return
    x.fullspan = 0
    for i = 1:length(x.args)
        x.fullspan += x.args[i].fullspan
    end
    x.span = x.fullspan - last(x.args).fullspan + last(x.args).span
    return
end

function Base.push!(e::EXPR, arg::EXPR)
    e.span = e.fullspan + arg.span
    e.fullspan += arg.fullspan
    setparent!(arg, e)
    push!(e.args, arg)
end

function Base.pushfirst!(e::EXPR, arg::EXPR)
    e.fullspan += arg.fullspan
    setparent!(arg, e)
    pushfirst!(e.args, arg)
end

function Base.pop!(e::EXPR)
    arg = pop!(e.args)
    e.fullspan -= arg.fullspan
    if isempty(e.args)
        e.span = 0
    else
        e.span = e.fullspan - last(e.args).fullspan + last(e.args).span
    end
    arg
end

function Base.append!(e::EXPR, args::Vector{EXPR})
    append!(e.args, args)
    for arg in args
        setparent!(arg, e)
    end
    update_span!(e)
end

function Base.append!(a::EXPR, b::EXPR)
    append!(a.args, b.args)
    for arg in b.args
        setparent!(arg, a)
    end
    a.fullspan += b.fullspan
    a.span = a.fullspan + last(b.span)
end

function INSTANCE(ps::ParseState)
    if isidentifier(ps.t)
        return mIDENTIFIER(ps)
    elseif isliteral(ps.t)
        return mLITERAL(ps)
    elseif iskw(ps.t)
        return mKEYWORD(ps)
    elseif isoperator(ps.t)
        return mOPERATOR(ps)
    elseif ispunctuation(ps.t)
        return mPUNCTUATION(ps)
    elseif kindof(ps.t) == Tokens.ERROR
        return EXPR(ErrorToken, nothing, ps.nt.startbyte - ps.t.startbyte, ps.t.endbyte - ps.t.startbyte + 1, val(ps.t, ps), NoKind, false, nothing, nothing, nothing, Unknown)
    else
        ps.errored = true
        return mErrorToken(Unknown)
    end
end

mutable struct File
    imports
    includes::Vector{Tuple{String,Any}}
    path::String
    ast::EXPR
    errors
end
File(path::String) = File([], [], path, EXPR(FileH, EXPR[]), [])

mutable struct Project
    path::String
    files::Vector{File}
end

function mUnaryOpCall(op::EXPR, arg::EXPR)
    fullspan = op.fullspan + arg.fullspan
    ex = EXPR(UnaryOpCall, EXPR[op, arg], fullspan, fullspan - arg.fullspan + arg.span)
    setparent!(op, ex)
    setparent!(op, ex)
    return ex
end
function mBinaryOpCall(arg1::EXPR, op::EXPR, arg2::EXPR)
    fullspan = arg1.fullspan + op.fullspan + arg2.fullspan
    ex = EXPR(BinaryOpCall, EXPR[arg1, op, arg2], fullspan, fullspan - arg2.fullspan + arg2.span)
    setparent!(arg1, ex)
    setparent!(op, ex)
    setparent!(arg2, ex)
    return ex
end
function mWhereOpCall(arg1::EXPR, op::EXPR, args::Vector{EXPR})
    ex = EXPR(WhereOpCall, EXPR[arg1; op; args], arg1.fullspan + op.fullspan, 0)
    setparent!(arg1, ex)
    setparent!(op, ex)
    for a in args
        ex.fullspan += a.fullspan
        setparent!(a, ex)
    end
    ex.span = ex.fullspan - last(args).fullspan + last(args).span
    return ex
end

mErrorToken(k::ErrorKind) = EXPR(ErrorToken, EXPR[], 0, 0, nothing, NoKind, false, nothing, nothing, nothing, k)
mErrorToken(x::EXPR, k) = EXPR(ErrorToken, EXPR[x], x.fullspan, x.span, nothing, NoKind, false, nothing, nothing, nothing, k)

TRUE() = mLITERAL(0, 0, "", Tokens.TRUE)
FALSE() = mLITERAL(0, 0, "", Tokens.FALSE)
NOTHING() = mLITERAL(0, 0, "", Tokens.NOTHING)
GlobalRefDOC() = EXPR(GlobalRefDoc, EXPR[])

typof(x::EXPR) = x.typ
valof(x::EXPR) = x.val
kindof(x::EXPR) = x.kind
kindof(t::Tokens.AbstractToken) = t.kind
parentof(x::EXPR) = x.parent
parentof(s::Scope) = s.parent
scopeof(x::EXPR) = x.scope
bindingof(x::EXPR) = x.binding
refof(x::EXPR) = x.ref

function setparent!(c, p)
    c.parent = p
    return c
end

function setscope!(x::EXPR, s = Scope())
    x.scope = s
    return x
end

function setref!(x::EXPR, r)
    x.ref = r
    return x
end

function setbinding!(x::EXPR)
    if typof(x) === TupleH
        for arg in x.args
            typof(arg) === PUNCTUATION && continue
            setbinding!(arg)
        end
    elseif typof(x) === Kw
        setbinding!(x.args[1], x)
    elseif typof(x) === Parameters
        for arg in x.args
            typof(arg) === PUNCTUATION && continue
            setbinding!(arg)
        end
    elseif typof(x) === InvisBrackets
        setbinding!(rem_invis(x))
    elseif typof(x) == UnaryOpCall && kindof(x.args[1]) === Tokens.DECLARATION
        return x
    else
        x.binding = Binding(x)
    end
    return x
end

function setbinding!(x::EXPR, binding)
    if typof(x) === TupleH
        for arg in x.args
            typof(arg) === PUNCTUATION && continue
            setbinding!(arg, binding)
        end
    elseif typof(x) === InvisBrackets
        setbinding!(rem_invis(x), binding)
    elseif typof(x) === IDENTIFIER || (typof(x) === BinaryOpCall && kindof(x.args[2]) === Tokens.DECLARATION)
        x.binding = Binding(str_value(get_name(x)), binding, nothing, [], nothing)
    end
    return x
end

function setiterbinding!(iter::EXPR)
    if typof(iter) === BinaryOpCall && kindof(iter.args[2]) in (Tokens.EQ, Tokens.IN, Tokens.ELEMENT_OF)
        setbinding!(iter.args[1], iter)
    end
    return iter
end

function mark_sig_args!(x::EXPR)
    if typof(x) === Call || typof(x) === TupleH
        if typof(x.args[1]) === InvisBrackets && typof(x.args[1].args[2]) === BinaryOpCall && kindof(x.args[1].args[2].args[2]) === Tokens.DECLARATION
            setbinding!(x.args[1].args[2])
        end
        for i = 2:length(x.args) - 1
            a = x.args[i]
            if typof(a) === Parameters
                for j = 1:length(a.args)
                    aa = a.args[j]
                    if !(typof(aa) === PUNCTUATION)
                        setbinding!(aa)
                    end
                end
            elseif !(typof(a) === PUNCTUATION)
                setbinding!(a)
            end
        end
    elseif typof(x) === WhereOpCall
        for i in 3:length(x.args)
            if !(typof(x.args[i]) === PUNCTUATION)
                setbinding!(x.args[i])
            end
        end
        mark_sig_args!(x.args[1])
    elseif typof(x) === BinaryOpCall
        if kindof(x.args[2]) == Tokens.DECLARATION
            mark_sig_args!(x.args[1])
        else
            setbinding!(x.args[1])
            setbinding!(x.args[3])
        end
    elseif typof(x) == UnaryOpCall && typof(x.args[2]) == InvisBrackets
        setbinding!(x.args[2].args[2])
    end
end
Base.getindex(x::EXPR, i) = x.args[i]

function strip_where_scopes(sig::EXPR)
    if typof(sig) === WhereOpCall
        setscope!(sig, nothing)
        strip_where_scopes(sig.args[1])
    end
end

function mark_typealias_bindings!(x::EXPR)
    x.binding = Binding(str_value(get_name(x.args[1])), x, nothing, [], nothing)
    setscope!(x)
    for i = 2:length(x.args[1].args)
        if typof(x.args[1].args[i]) === IDENTIFIER
            setbinding!(x.args[1].args[i])
        elseif typof(x.args[1].args[i]) === BinaryOpCall && kindof(x.args[1].args[i].args[2]) === Tokens.ISSUBTYPE && typof(x.args[1].args[i].args[1]) === IDENTIFIER
            setbinding!(x.args[1].args[i].args[1])
        end
    end
    return x
end


text:
@testset "Random" begin
    @testset for k in 1:32
        for _ in 1:10
            kmer = rand(DNAKmer{k})
            @test isa(kmer, DNAKmer{k})
            kmer = rand(RNAKmer{k})
            @test isa(kmer, RNAKmer{k})
        end

        for size in [0, 1, 2, 5, 10, 100]
            @test length(rand(DNAKmer{k}, size)) == size
            @test length(rand(RNAKmer{k}, size)) == size
        end

        kmers = rand(DNAKmer{k}, 10_000)
        for i in 1:k
            a = sum([kmer[i] for kmer in kmers] .== DNA_A)
            c = sum([kmer[i] for kmer in kmers] .== DNA_C)
            g = sum([kmer[i] for kmer in kmers] .== DNA_G)
            t = sum([kmer[i] for kmer in kmers] .== DNA_T)
            @test 2200 ≤ a ≤ 2800
            @test 2200 ≤ c ≤ 2800
            @test 2200 ≤ g ≤ 2800
            @test 2200 ≤ t ≤ 2800
            @test a + c + g + t == 10_000
        end
    end
end


text:
# Since Zygote is at its early stage
# this patch contains some workaround of known issues
# it should be able to get removed once those issues are fixed

using Zygote, LuxurySparse, Yao, YaoBase, SparseArrays, BitBasis

# LuxurySparse constructor adjoints
Zygote.@adjoint (::Type{T})(perms, vals) where T <: PermMatrix = T(perms, vals), Δ -> nothing
Zygote.@adjoint (::Type{T})() where T <: IMatrix = T(), Δ -> nothing
Zygote.@adjoint Base.:(*)(A::Number, B::PermMatrix) = A * B, Δ->(sum(Δ .* B), A * Δ)
Zygote.@adjoint Base.:(*)(A::PermMatrix, B::Number) = A * B, Δ->(A * Δ, sum(Δ .* B))

Zygote.@adjoint SparseArrays.SparseMatrixCSC(A::PermMatrix) = SparseMatrixCSC(A), Δ->(Δ, )

# BitBasis constructor (no gradient needed)
Zygote.@adjoint BitBasis.onehot(::Type{T}, nbits::Int, x::Integer, nbatch::Int) where T = onehot(T, nbits, x, nbatch), Δ -> nothing

# state(r::ArrayReg)
Zygote.@adjoint function Base.getfield(r::ArrayReg, n::Symbol)
    getfield(r, n), Δ -> (ArrayReg(Δ), nothing)
end

# workaround of upstream
Zygote.@adjoint Base.:(-)(a, b) = a-b, Δ -> (Δ, -Δ)
Zygote.@adjoint Base.:(+)(a, b) = a+b, Δ -> (Δ, Δ)

# require mutation support
Zygote.@adjoint! function copyto!(xs::AbstractVector, ys::Tuple)
    xs_ = copy(xs)
    copyto!(xs, ys), function (dxs)
        copyto!(xs_, xs)
        return (nothing, Tuple(dxs))
    end
end

max_stars_repo_name: MikeInnes/zygote-paper

max_issues_repo_name: MikeInnes/zygote-paper

max_forks_repo_name: MikeInnes/zygote-paper


text:
# MIT license
# Copyright (c) Microsoft Corporation. All rights reserved.
# See LICENSE in the project root for full license information.

const sin60 = .5*sqrt(3)
const cos60 = .5
"""coordinates of a unit hexagon tile polygon with unit length sides centered about the point (0,0)"""
const hexcoords = [
		1 0;
		cos60 -sin60;
		-cos60 -sin60;
		-1 0;
		-cos60 sin60;
		cos60 sin60
		]

struct HexBasis1{N,T} <: Basis{N,T}
    HexBasis1(::Type{T} = Float64) where{T<:Real} = new{2,T}()
end
export HexBasis1

basis(::HexBasis1{2,T}) where{T} = SMatrix{2,2,T}(T(1.5),T(.5)*sqrt(T(3)),T(1.5),T(-.5)*(sqrt(T(3))))

# SVector{2,SVector{2,T}}(hexe₁(T),hexe₂(T))

"""Returns the vertices of the unit tile polygon for the basis"""
function tilevertices(::HexBasis1{2,T}) where{T}
sin60 = T(.5)*sqrt(T(3))
cos60 = T(.5)
return SMatrix{2,6}(
		1, 0,
		cos60, -sin60,
		-cos60, -sin60,
		-1, 0,
		-cos60, sin60,
		cos60, sin60)
end
export tilevertices

# coordinate offsets to move up, down, etc., numsteps hex cells in a hex lattice defined in HexBasis1
hexup(numsteps = 1) = numsteps .* (1,-1)
hexdown(numsteps = 1) = numsteps .* (-1,1)
hexupright(numsteps = 1) = numsteps .* (1,0)
hexdownleft(numsteps = 1) = numsteps .* (-1,0)
hexdownright(numsteps = 1) = numsteps .* (0,1)
hexupleft(numsteps = 1) = numsteps .* (0,-1)

""" (i,j) offsets of a ring represented in the HexBasis1 lattice basis. The hexagonal grid is assumed to start from the bottom center hex cell. To generate larger rings repeat this pattern
 ring 1 = (1, 0)      (1, -1)       (0, -1)       (-1, 0)       (-1, 1)       (0, 1)
 ring 2 = (1, 0)(1, 0)(1, -1)(1, -1)(0, -1)(0, -1)(-1, 0)(-1, 0)(-1, 1)(-1, 1)(0, 1)(0, 1)
"""
const hexcycle = SVector{6,NTuple{2,Int64}}(hexupright(),hexup(),hexupleft(),hexdownleft(),hexdown(),hexdownright())

"""returns (i,j) offsets of ring n defined in the HexBasis1 lattice basis."""
ringoffsets(n::Int64) = ringoffsets(Val{n})
function ringoffsets(::Type{Val{N}}) where N
    temp = MVector{N*6,NTuple{2,Int64}}(undef)

    for i in 1:6
        for j in 1:N
            temp[(i-1)*N + j] = hexcycle[i]
        end
    end

    return SVector{N*6,NTuple{2,Int64}}(temp)
end
export ringoffsets

"""Returns ```centerpoint``` plus all its neighbors in a ring of size n. This is a hexagonal region of the lattice n times as large as the unit hexagon cell.

Example:

```
Vis.@wrapluxor Vis.drawhexcells(50, hexregion((0,0),2))
```
"""
function region(::Type{HexBasis1},centerpoint::Tuple{Int64,Int64}, n::Int64)
    f(i) = i==0 ? () : (neighbors(HexBasis1,centerpoint,i)...,f(i-1)...)
    cells = (centerpoint...,f(n)...)
    numcells = length(cells) ÷ 2
    println(numcells)
    # convert to matrix form since that is what the drawing code expects
    temp = MMatrix{2,numcells,Int64}(undef)
    for i in 1:numcells
        println(cells[i])
        temp[1,i] = cells[2*i-1]
        temp[2,i] = cells[2*i]
    end

    return SMatrix{2,numcells,Int64}(temp)
end
export region

"""Returns all hex cells contained in the ring of size n centered around centerpoint. Use region if you want all the cells contained in the ring of size n."""
function neighbors(::Type{HexBasis1}, centerpoint::Tuple{Int64,Int64},n::Int) where{T}
    temp = MMatrix{2,n*6,Int64}(undef)
    hoffsets = ringoffsets(Val{n})
    latticepoint = centerpoint .+ n .* (hexdown())

    #convert hexoffsets to absolute coordinate positions
    for i in 1:size(temp)[2]
        temp[1,i] = latticepoint[1]
        temp[2,i] = latticepoint[2]
        latticepoint = latticepoint .+ hoffsets[i] #last computed value of latticepoint won't be used
    end
    return SMatrix{2,n*6,Int64}(temp)
end
export neighbors

function xbounds(numi)
    numevens = div(numi,2)
    numodds = numevens + mod(numi,2)
    basewidth = numodds + 2*numevens + 1

    if iseven(numi)
        maxx = basewidth
    else
        maxx = basewidth + sin(deg2rad(30))
    end

    minx  = -maxx
    return (minx,maxx)
end

function ybounds(numj)
    maxy = 2*sin60*(numj + 1)
    return (-maxy,maxy-sin60)
end

"""computes bounding box of hex tiles arranged in an approximately rectangular layout, with 2*numj+1 hexagons in a row, and 2*numi+1 hexagons in a column"""
function bbox(numi,numj)
    return (xbounds(numi),ybounds(numj))

end
export bbox

"""computes the basis coordinates of hexagonal cells defined in the Hex1Basis in a rectangular pattern 2*numi+1 by 2*numj+1 wide and high respectively."""
function hexcellsinbox(numi,numj)
    #i2 and j2 are coordinates in the basis (1.5,.5√3),(0,-√3). The conditional tests are simpler in this basis. Convert to Hex1Basis1 coordinates before returning.
    hexbasis2to1(i,j) = [i+j,-j]

    result = Matrix{Int}(undef,2,(2*numi+1)*(2*numj+1))
    for i2 in -numi:numi
        let offsetj
            if i2 < 0
            offsetj = -abs(div(i2,2))
            else
                offsetj = div(i2+1,2)
            end
            for j2 in -numj:numj
                # conversion from Hex2Basis i2,j2 to Hex1Basis i,j is
                # (i,j) = (i2,0) + j2*(1,-1)
                jtemp = j2-offsetj
                # i1 = i2 + jtemp
                # j1 = -jtemp
                result[:,(i2+numi)*(2*numj+1) + j2+numj+1] = hexbasis2to1(i2,jtemp)
            end
        end
    end
    return result
end
export Repeat

struct HexBasis3{N,T}<:Basis{N,T}
    HexBasis3(::Type{T} = Float64) where{T} = new{2,T}()
end
export HexBasis3

function tilevertices(::HexBasis3{2,T}) where{T}
    sin60 = T(.5)*sqrt(T(3))
    cos60 = T(.5)
    return SMatrix{2,6,T}(
        0, 1,
        -sin60, cos60,
        -sin60, -cos60,
        0, -1,
        sin60, -cos60,
        sin60, cos60)
end

basis(::HexBasis3{2,T}) where{T} = SMatrix{2,2,T}(T(2*sin60),T(0),T(sin60),T(1.5))


text:
mutable struct TrackedReal{T<:Real}
    data::T
    grad::Union{Nothing,T}
    children::Dict
    # this field is only need for printing the graph. you can safely remove it.
    name::String
end

track(x::Real,name="") = TrackedReal(x,nothing,Dict(),name)

function Base.show(io::IO, x::TrackedReal)
    t = isempty(x.name) ? "(tracked)" : "(tracked $(x.name))"
    print(io, "$(x.data) $t")
end

function accum!(x::TrackedReal)
    if isnothing(x.grad)
        x.grad = sum(accum!(v)*w for (v,w) in x.children)
    end
    x.grad
end

function gradient(f, args::Real...)
    ts = track.(args)
    y  = f(ts...)
    y.grad = 1.0
    accum!.(ts)
end

##########  RULES  #############################################################

function Base.:*(a::TrackedReal, b::TrackedReal)
    z = track(a.data * b.data, "*")
    a.children[z] = b.data  # dz/da=b
    b.children[z] = a.data  # dz/db=a
    z
end
function Base.:+(a::TrackedReal{T}, b::TrackedReal{T}) where T
    z = track(a.data + b.data, "+")
    a.children[z] = one(T)
    b.children[z] = one(T)
    z
end
function Base.sin(x::TrackedReal)
    z = track(sin(x.data), "sin")
    x.children[z] = cos(x.data)
    z
end

##########  Optimizion 2D function  ############################################

using Plots
g(x,y) = y*y + sin(x)
cscheme = cgrad(:RdYlBu_5, rev=true)
p1 = contour(-4:0.1:4, -2:0.1:2, g, fill=true, c=cscheme, xlabel="x", ylabel="y")
display(p1)

function descend(f::Function, λ::Real, args::Real...)
    Δargs = gradient(f, args...)
    args .- λ .* Δargs
end

function minimize(f::Function, args::T...; niters=20, λ=0.01) where T<:Real
    paths = ntuple(_->Vector{T}(undef,niters), length(args))
    for i in 1:niters
        args = descend(f, λ, args...)
        @info f(args...)
        for j in 1:length(args)
            paths[j][i] = args[j]
        end
    end
    paths
end

xs1, ys1 = minimize(g, 1.5, -2.4, λ=0.2, niters=34)
xs2, ys2 = minimize(g, 1.8, -2.4, λ=0.2, niters=16)

scatter!(p1, [xs1[1]], [ys1[1]], markercolor=:black, marker=:star, ms=7, label="Minimum")
scatter!(p1, [xs2[1]], [ys2[1]], markercolor=:black, marker=:star, ms=7, label=false)
scatter!(p1, [-π/2], [0], markercolor=:red, marker=:star, ms=7, label="Initial Point")
scatter!(p1, xs1[1:1], ys1[1:1], markercolor=:black, label="GD Path", xlims=(-4,4), ylims=(-2,2))

anim = @animate for i in 1:max(length(xs1), length(xs2))
    if i <= length(xs1)
        scatter!(p1, xs1[1:i], ys1[1:i], mc=:black, lw=3, xlims=(-4,4), ylims=(-2,2), label=false)
    end
    if i <= length(xs2)
        scatter!(p1, xs2[1:i], ys2[1:i], mc=:black, lw=3, label=false)
    end
    p1
end

gif(anim, "gd-path.gif", fps=15)


text:
function rat(x::AbstractFloat, tol::Real=eps(x))::Rational
    p, q, pp, qq  = copysign(1,x), 0, 0, 1
    x, y = abs(x), 1.0
    r, a = modf(x)
    nt, t, tt = tol, 0.0, tol

    while r > nt        # convergents of the continued fraction: np//nq = (p*a + pp) // (q*a + qq)
        np, nq = Int(a).*(p,q) .+ (pp,qq)
        p, pp, q, qq = np, p, nq, q

        x, y = y, r     # instead of the inexact 1/r...
        a, r = divrem(x,y)

        t, tt = nt, t   # maintain x = (p + (-1)^i * r) / q
        nt = a*t+tt
    end

    i = Int(cld(x-tt,y+t))   # find optimal semiconvergent: smallest i such that x-i*y < i*t+tt
    return (i*p+pp) // (i*q+qq)
end

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
@testset "QuadrotorGame" begin

	# Test QuadrotorGame
	p = 3
	model = QuadrotorGame(p=p)
	@test typeof(model) <: QuadrotorGame{model.n, model.m, model.p}
	@test model.n == 12p
	@test model.m == 4p
	@test model.p == p
	@test model.ni == 12*ones(Int,p)
	@test model.mi == 4*ones(Int,p)
	@test model.pu[1] == SVector{4,Int}((1,4,7,10))
	@test model.pu[2] == SVector{4,Int}((2,5,8,11))
	@test model.pu[3] == SVector{4,Int}((3,6,9,12))

	@test model.px[1] == SVector{2,Int}((1,4))
	@test model.px[2] == SVector{2,Int}((2,5))
	@test model.px[3] == SVector{2,Int}((3,6))

	@test model.pz[1] == SVector{12,Int}((1,4,7,10,13,16,19,22,25,28,31,34))
	@test model.pz[2] == SVector{12,Int}((2,5,8,11,14,17,20,23,26,29,32,35))
	@test model.pz[3] == SVector{12,Int}((3,6,9,12,15,18,21,24,27,30,33,36))

	@test size(model) == (12p, 4p, model.pu, p)

	T = Float64
	n = 12
	m = 4

	p = 3
	model = QuadrotorGame(p=p)
	x = rand(SVector{n*p,T})
	u = rand(SVector{m*p,T})
	@test (@ballocated $(Algames.forces)($model, $x, $u, $1)) == 0
	@test (@ballocated $(Algames.moments)($model, $x, $u, $2)) == 0
	@test (@ballocated $(Algames.wrenches)($model, $x, $u, $2)) == 0
	@test (@ballocated RobotDynamics.dynamics($model, $x, $u, $1)) == 0
	for p = 1:4
		model = QuadrotorGame(p=p)
		x = rand(SVector{n*p,T})
		u = rand(SVector{m*p,T})
		@test (@ballocated RobotDynamics.dynamics($model, $x, $u)) == 0
	end

end


text:
"""
    SGD(::Vector{XVariable}; lr=1e-4, L1decay=0.001, L2decay=0.01)

Implements stochastic gradient descent
"""
mutable struct SGD <: Optimizer
    xparams::Vector{XVariable}
    lr::AbstractFloat
    L1decay::AbstractFloat
    L2decay::AbstractFloat
    name::String
    function SGD(xparams::Vector{XVariable}; lr=1e-4, L1decay=0.001, L2decay=0.01)
        new(xparams, lr, L1decay, L2decay, "SGD")
    end
end

# pretty printing
function Base.show(io::IO, O::SGD)
    print("SGD(lr=$(O.lr), L1decay=$(O.L1decay), L2decay=$(O.L2decay))")
end

function update!(O::SGD; clipfn::Function=LPInfNormClip, clipvalue=10.0)
    μ = - O.lr
    λ₁ = O.L1decay
    λ₂ = O.L2decay

    for i = 1:length(O.xparams)
        c , θ = O.xparams[i]
        ∇ = clipfn(setNanInfZero(δ(θ)), clipvalue)
        𝒗 = ᵛ(θ)
        if c == 'w'
            if λ₁==0 && λ₂==0
                @. 𝒗 += μ * ∇
            elseif λ₁==0 && λ₂!=0
                @. 𝒗 += μ * (∇ + λ₂ * 𝒗)
            elseif λ₁!=0 && λ₂==0
                @. 𝒗 += μ * (∇ + λ₁ * sign(𝒗))
            else  # λ₁!=0 && λ₂!=0
                @. 𝒗 += μ * (∇ + λ₁ * sign(𝒗) + λ₂ * 𝒗)
            end
        else
            @. 𝒗 += μ * ∇
        end
    end
end

max_stars_repo_name: sonosole/Delta

max_issues_repo_name: sonosole/Delta

max_forks_repo_name: sonosole/Delta


text:
using MetaGraphs: AbstractMetaGraph

## Linear algebra API for AbstractMetaGraph

function degrees(mg::AbstractMetaGraph, T::DataType=eltype(mg); dir::Symbol=:out)
    degrees(adjacency_matrix(mg.graph, T; dir=dir), T; dir=dir)
end

function degree_matrix(mg::AbstractMetaGraph, T::DataType=eltype(mg); dir::Symbol=:out)
    degree_matrix(adjacency_matrix(mg.graph, T; dir=dir), T; dir=dir)
end

function inv_sqrt_degree_matrix(mg::AbstractMetaGraph, T::DataType=eltype(mg); dir::Symbol=:out)
    inv_sqrt_degree_matrix(adjacency_matrix(mg.graph, T; dir=dir), T; dir=dir)
end

function laplacian_matrix(mg::AbstractMetaGraph, T::DataType=eltype(mg); dir::Symbol=:out)
    laplacian_matrix(adjacency_matrix(mg.graph, T; dir=dir), T; dir=dir)
end

function normalized_laplacian(mg::AbstractMetaGraph, T::DataType=eltype(mg); selfloop::Bool=false)
    adj = adjacency_matrix(mg.graph, T)
    selfloop && (adj += I)
    normalized_laplacian(adj, T)
end

## Convolution layers accepting AbstractMetaGraph

GCNConv(g::AbstractMetaGraph, ch::Pair{<:Integer,<:Integer}, σ=identity; kwargs...) =
    GCNConv(g.graph, ch, σ; kwargs...)

ChebConv(g::AbstractMetaGraph, ch::Pair{<:Integer,<:Integer}, k::Integer; kwargs...) =
    ChebConv(g.graph, ch, k; kwargs...)

GraphConv(g::AbstractMetaGraph, ch::Pair{<:Integer,<:Integer}, aggr=:add; kwargs...) =
    GraphConv(g.graph, ch, aggr; kwargs...)

GATConv(g::AbstractMetaGraph, ch::Pair{<:Integer,<:Integer}; kwargs...) =
    GATConv(g.graph, ch; kwargs...)

GatedGraphConv(g::AbstractMetaGraph, out_ch::Integer, num_layers::Integer; kwargs...) =
    GatedGraphConv(g.graph, out_ch, num_layers; kwargs...)

EdgeConv(g::AbstractMetaGraph, nn; kwargs...) = EdgeConv(g.graph, nn; kwargs...)


text:
using Documenter, Hecke

function Markdown.plain(io::IO, ::Markdown.HorizontalRule)
           println(io, "-"^3)
end

makedocs(
    modules = Hecke,
    clean   = true,
    doctest = false,
)

# Hack around to get syntax highlighting working
cd(Pkg.dir("Hecke", "docs"))
cp("application-f78e5cb881.palette.css", "build/application-f78e5cb881.palette.css", remove_destination = true)
cp("application-e2807e330f.css", "build/application-e2807e330f.css", remove_destination = true)

deploydocs(
    deps = Deps.pip("pygments",
                    "mkdocs==0.16.3",
                    "python-markdown-math",
                    "mkdocs-cinder"),
    repo = "github.com/thofma/Hecke.jl.git",
    osname = "linux",
    julia = "0.6",
)


text:
arg_dict = Dict{Any, Any}(
"physics" => "Euler",
"dimensions" => 2,
"use_DG" => true,
"operator_type" => "SBPOmega",
"run_type" => 40,
"jac_type" => 3,
"jac_method" => 2,
"real_time" => false,
#"jac_type" => 3,
#"jac_method" => 2,
"order" => 1,
"Flux_name" => "RoeFlux",
#"reordering_algorithm" => "adjacency",
#"reordering_start_coords" => [50., 0],
#"Flux_name" => "IRFlux",
"IC_name" => "ICFreeStream",
#"volume_integral_type" => 2,
#"face_integral_type" => 2,
#"Volume_flux_name" => "IRFlux",
#"FaceElementIntegral_name" => "ESLW2FaceIntegral",
"numBC" => 2,
"BC1" => [5],  # surface of airfoil, symmetry plane
"BC1_name" => "noPenetrationBC",
"BC2" => [8],
"BC2_name" => "FreeStreamBC",
# set parameters
"Ma" => 0.2,
"aoa" => 2,  # degrees
#"delta_t" => 0.0005,
#"CFL" => 0.5,
"t_max" => 200000000000.0,
"smb_name" => "SRCMESHES/0012_small.smb",
"dmg_name" => "SRCMESHES/0012_small.dmg",
"itermax" => 30,
"output_freq" => 1,
"write_vis" => true,
"res_abstol" => 1e-7,
"res_reltol" => 1e-50,
"krylov_reltol" => 1e-4,
"krylov_abstol" => 1e-6,
"do_postproc" => false,
"solve" => true,
#"write_entropy" => true,
#"write_entropy_freq" => 1,
"newton_verbosity" => 4,
#"calc_jac_explicit" => true,
)


text:
module G
function f(q::Vector{T}, q_dot::Vector{T}, xi::T) where T
    l1 = q[1]
    l2 = q[2]
    l3 = q[3]
    l1_dot = q_dot[1]
    l2_dot = q_dot[2]
    l3_dot = q_dot[3]

    [110   0   0;
0 110   0;
0   0 110]
end
end


text:
header("FLATTEN")

@testset "deref" begin
    str = """
    <?xml version="1.0"?>
    <pnml xmlns="http://www.pnml.org/version-2009/grammar/pnml">
        <net id="net0" type="pnmlcore">
            <page id="page1">
                <place id="p1"/>
                <transition id ="t1"/>
                <arc id="a1" source="p1" target="t1"/>
                <arc id="a12" source="t1" target="rp1"/>
                <referencePlace id="rp1" ref="p2"/>
            </page>
            <page id="page2">
                <place id="p2"/>
                <transition id ="t2"/>
                <arc id="a2" source="t2" target="p2"/>
                <arc id="a22" source="t2" target="rp2"/>
                <referencePlace id="rp2" ref="p3"/>
                <referenceTransition id="rt2" ref="t3"/>
            </page>
            <page id="page3">
                <place id="p3"/>
                <transition id ="t3"/>
                <arc id="a3" source="t3" target="p3"/>
            </page>
        </net>
    </pnml>
    """
    model = parse_str(str)
    net = PNML.first_net(model)
    @test_call PNML.first_net(model)

    PNML.flatten_pages!(net)
    @test_call PNML.flatten_pages!(net)
    @test typeof(net) <: PNML.PnmlNet
    #@show net
end


text:
export get_E_all

#ENERGIES
@inline function get_E_all(s::Simulation)::MVector{3, Float64}
    #Units are STILL not clear
    E_au_au = get_V_au_au(s)
    E_coup = get_E_coup(s)
    E_neutral = get_E_neutral(s) + E_au_au
    E_ion = get_E_ion(s) + E_au_au
    hp = MVector{3, Float64}(E_neutral, E_ion, E_coup)
    return hp
end

#---------------------------Coupling----------------------------------

function get_E_coup(s::Simulation)::Float64 #check
    E_au_n_coup = get_E_au_n_coup(s)
    E_au_o_coup = get_E_au_o_coup(s)
    return E_au_n_coup + E_au_o_coup
end

#---------------------------Au-N Coupling----------------------------------
@inline function get_E_au_n_coup_loop!(x, V, i, delta_xn, cutoff_test,
    bool_cutoff_test, s::Simulation)::Float64
    @views delta_xn =  x[i+2, :] - x[1, :]
    delta_xn[1] = @.delta_xn[1] - cell[1]*round(delta_xn[1]/cell[1])
    delta_xn[2] = @.delta_xn[2] - cell[2]*round(delta_xn[2]/cell[2])

    cutoff_test .= abs.(delta_xn) .- cutoff
    bool_cutoff_test .= cutoff_test .<= 0.0
    sum_bool = sum(bool_cutoff_test)
    if sum_bool != zero(sum_bool)
        norm_delta_xn = norm(delta_xn)
        if norm_delta_xn <= cutoff
            V_iterate = coup_a_N / (1.0 + coup_b_N * exp(coup_β_N*norm_delta_xn))
            V_cut = coup_cutoff_N
            V = V + V_iterate - V_cut
        end
    end
    return V
end

@inbounds @inline function get_E_au_n_coup(s::Simulation)::Float64
    V = 0.0
    delta_xn = Vector{Float64}(undef, 3)
    cutoff_test = Vector{Float64}(undef, 3)
    bool_cutoff_test = BitArray{1}(undef, 3)
    for i in 1:N
        V = get_E_au_n_coup_loop!(s.x, V, i, delta_xn, cutoff_test, bool_cutoff_test, s)
    end
    return V
end
#---------------------------Au-O Coupling----------------------------------
@inline function get_E_au_o_coup_loop!(x, V, i, delta_xo, cutoff_test,
     bool_cutoff_test, s::Simulation)::Float64
    @views delta_xo = x[i+2, :] .- x[2, :]
    delta_xo[1] = delta_xo[1] - cell[1]*round(delta_xo[1]/cell[1])
    delta_xo[2] = delta_xo[2] - cell[2]*round(delta_xo[2]/cell[2])

    cutoff_test .= abs.(delta_xo) .- cutoff
    bool_cutoff_test .= cutoff_test .<= 0.0
    sum_bool = sum(bool_cutoff_test)
    if sum_bool != zero(sum_bool)
        norm_delta_xo = norm(delta_xo)
        if norm_delta_xo <= cutoff
            V_coup = coup_a_O /(1.0 + coup_b_O * exp(coup_β_O*norm_delta_xo))
            V_cut = coup_cutoff_O
            V = V + V_coup - V_cut
        end
    end
    return V
end

@inline @inbounds function get_E_au_o_coup(s::Simulation)::Float64
    V = 0.0
    delta_xo = Vector{Float64}(undef, 3)
    cutoff_test = Vector{Float64}(undef, 3)
    bool_cutoff_test = BitArray{1}(undef, 3)
    for i in 1:N
        V = get_E_au_o_coup_loop!(s.x, V, i, delta_xo, cutoff_test, bool_cutoff_test, s)
    end
    return V
end
#---------------------------Ionic----------------------------------

function get_E_ion(s::Simulation)::Float64  #check
    E_n_o_ion= get_E_n_o_ion(s)
    E_au_n_ion = get_E_au_n_ion(s)
    E_au_o_ion = get_E_au_o_ion(s)
    E_image = get_E_image_ionic(s)
    return E_n_o_ion + E_au_n_ion+ E_au_o_ion + E_image + KI
end

#---------------------------Ionic N O ----------------------------------

function get_E_n_o_ion(s::Simulation)::Float64
    norm_r_no = norm(s.Δ_no)
    return F_i * (1 - exp(-δ_i*(norm_r_no - rNO_e_i)))^2
end

#---------------------------Ionic Au O ----------------------------------
@inline function get_E_au_o_ion_loop(x::Array{Float64,2}, V::Float64, i::Int,
    delta_xo, cutoff_test, bool_cutoff_test, s::Simulation)::Float64
    @views delta_xo .=  x[i+2, :] .- x[2, :]
    delta_xo[1] = delta_xo[1] - cell[1]*round(delta_xo[1]/cell[1])
    delta_xo[2] = delta_xo[2] - cell[2]*round(delta_xo[2]/cell[2])

    cutoff_test .= abs.(delta_xo) .- cutoff
    bool_cutoff_test .= cutoff_test .<= 0.0
    sum_bool = sum(bool_cutoff_test)
    if sum_bool != zero(sum_bool)
        norm_delta_xo = norm(delta_xo)
        if norm_delta_xo <= cutoff
            V_iterate = A_n * (
            exp(-α_i * norm_delta_xo)
            - exp_alpha_i_cutoff
            )
            V = V + V_iterate
        end
    end
    return V
end

@inline @inbounds function get_E_au_o_ion(s::Simulation)::Float64
    V = 0.0
    delta_xo = Vector{Float64}(undef, 3)
    cutoff_test = Vector{Float64}(undef, 3)
    bool_cutoff_test = BitArray{1}(undef, 3)
    for i in 1:N
        V = get_E_au_o_ion_loop(s.x, V, i, delta_xo, cutoff_test, bool_cutoff_test, s)
    end
    return V
end

#---------------------------Ionic Au N ----------------------------------
@inline @inbounds function get_E_au_n_ion_loop(x::Array{Float64,2}, V::Float64,
     i::Int, delta_xn, cutoff_test, bool_cutoff_test, s::Simulation)::Float64
    @views delta_xn .=  x[i+2, :] .- x[1, :]
    delta_xn[1] = delta_xn[1] - cell[1]*round(delta_xn[1]/cell[1])
    delta_xn[2] = delta_xn[2] - cell[2]*round(delta_xn[2]/cell[2])

    cutoff_test .=  abs.(delta_xn) .- cutoff
    bool_cutoff_test .=  cutoff_test .<= 0.0
    sum_bool = sum(bool_cutoff_test)
    if sum_bool != zero(sum_bool)
        norm_delta_xn = norm(delta_xn)
        if norm_delta_xn <= cutoff
            norm_Δ_no = norm(s.Δ_no)
            cos_theta = (s.Δ_no[3])/norm_Δ_no
            V_iterate =  B_i * (
            exp(-2*β_i *(norm_delta_xn - rN_sur_e_i)) -
            exp_beta_i_cutoff1 -
            2*cos_theta^2 * exp(-β_i *(norm_delta_xn - rN_sur_e_i))
            - exp_beta_i_cutoff2
            )
            V = V + V_iterate
        end
    end
    return V
end

@inline function get_E_au_n_ion(s::Simulation)::Float64
    V = 0.0
    delta_xn = Vector{Float64}(undef, 3)
    cutoff_test = Vector{Float64}(undef, 3)
    bool_cutoff_test = BitArray{1}(undef, 3)
    for i in 1:N
        V = get_E_au_n_ion_loop(s.x, V, i, delta_xn, cutoff_test, bool_cutoff_test, s)
    end
    return V
end

#---------------------------Ionic Image Potential----------------------------------
@inline function get_E_image_ionic(s::Simulation)::Float64
    @views zcom = zcom_NO(s.x[1:2, :])
    return -D_i/sqrt(C_i^2 + ((zcom - z0_i))^2)
end
#---------------------------Ionic N-O----------------------------------
@inline function get_E_n_o_ion(s::Simulation)::Float64
    @views r_no = s.x[1, :] - s.x[2,:]
    norm_r_no = norm(r_no)
    return F_i * (1 - exp(-δ_i*(norm_r_no - rNO_e_i)))^2
end
#---------------------------Neutral----------------------------------

@inline function get_E_neutral(s::Simulation)::Float64  #check
    E_n_o_neutral = get_E_n_o_neutral(s)
    E_au_n_neutral = get_E_au_n_neutral(s)
    E_au_o_neutral = get_E_au_o_neutral(s)
    return E_n_o_neutral + E_au_n_neutral + E_au_o_neutral
end

#---------------------------Neutral N-O----------------------------------

@inline function get_E_n_o_neutral(s::Simulation)::Float64 #its good
    norm_r_no = norm(s.Δ_no)
    return F_n * (1 - exp(-δ_n*(norm_r_no - r_0_NO)))^2
end

#---------------------------Neutral AU_O----------------------------------
@inline function get_E_au_o_neutral_loop!(x, i, V, delta_xn, cutoff_test,
     bool_cutoff_test, s::Simulation)::Float64 #i think its ok, small numerical instability
    @views delta_xn .= x[i+2, :] .- x[2, :]
    delta_xn[1] = delta_xn[1] - cell[1]*round(delta_xn[1]/cell[1])
    delta_xn[2] = delta_xn[2] - cell[2]*round(delta_xn[2]/cell[2])

    cutoff_test .= abs.(delta_xn) .- cutoff
    bool_cutoff_test .= cutoff_test .<= 0.0
    sum_bool = sum(bool_cutoff_test)
    if sum_bool != zero(sum_bool)
        norm_delta_xn = norm(delta_xn)
        if norm_delta_xn <= cutoff
            V = V + A_n * (exp(-α_n*norm_delta_xn) - exp_alpha_n_cutoff)
        end
    end
    return V
end

@inbounds @inline function get_E_au_o_neutral(s::Simulation)::Float64
    V = 0.0
    delta_xn = zeros(Float64, 3)
    cutoff_test = zeros(Float64, 3)
    bool_cutoff_test = zeros(Bool, 3)
    for i in 1:N
        V = get_E_au_o_neutral_loop!(s.x, i, V, delta_xn, cutoff_test, bool_cutoff_test, s)
    end
    return V
end
#---------------------------Neutral AU_N----------------------------------

@inline function get_E_au_n_neutral_loop(x::Array{Float64,2}, i::Int, V::Float64,
     delta_xn, cutoff_test, bool_cutoff_test, s::Simulation)::Float64
    @views delta_xn .=  x[i+2, :] .- x[1, :]
    delta_xn[1] = delta_xn[1] - cell[1]*round(delta_xn[1]/cell[1])
    delta_xn[2] = delta_xn[2] - cell[2]*round(delta_xn[2]/cell[2])

    cutoff_test .= abs.(delta_xn) .- cutoff
    bool_cutoff_test .=  cutoff_test .<= 0.0
    sum_bool = sum(bool_cutoff_test)
    if sum_bool != zero(sum_bool)
        norm_delta_xn = norm(delta_xn)
        if norm_delta_xn <= cutoff
            V = V + B_n * (exp(-β_n*norm_delta_xn) - exp_beta_n_cutoff)
        end
    end
    return V
end

@inbounds @inline function get_E_au_n_neutral(s::Simulation)::Float64
    V = 0.0
    delta_xn = Vector{Float64}(undef, 3)
    cutoff_test = Vector{Float64}(undef, 3)
    bool_cutoff_test = BitArray{1}(undef, 3)
    for i in 1:N
        V = get_E_au_n_neutral_loop(s.x, i, V, delta_xn, cutoff_test, bool_cutoff_test, s)
    end
    return V
end

#---------------------------GOLD AU AU LATTICE----------------------------------
@inline function V_au_au_loop_if(x::float_array, V::Float64, i::Int, j::Int,
     xm, xi, temp,r, qf_1, qf_2, s::Simulation)::Float64
    if s.nn_arr[i, j] != 0
        m = s.nn_arr[i, j] + 2
        @views xm .= x[m, :]
        @views xi .= x[i+2, :]
        temp .= xm .- xi
        temp .= temp./cell
        temp .= temp .- floor.(temp .+ 0.5)
        @views r .=  temp .* cell .- r0[:, j]
        if j == 1 || j== 7
            mul!(qf_2, d1_new, r)
            qf_1 = dot(r, qf_2)
            V = V + qf_1
        elseif j == 2 || j== 8
            mul!(qf_2, d2_new, r)
            qf_1 = dot(r, qf_2)
            V = V + qf_1
        elseif j == 3 || j== 9
            mul!(qf_2, d3_new, r)
            qf_1 = dot(r, qf_2)
            V = V + qf_1
        elseif j == 4 || j== 10
            mul!(qf_2, d4_new, r)
            qf_1 = dot(r, qf_2)
            V = V + qf_1
        elseif j == 5|| j== 11
            mul!(qf_2, d5_new, r)
            qf_1 = dot(r, qf_2)
            V = V + qf_1
        elseif j == 6 || j== 12
            mul!(qf_2, d6_new, r)
            qf_1 = dot(r, qf_2)
            V = V + qf_1
        end
    end
    return V
end
@inline @inbounds function V_au_au_loop(x::float_array, V::Float64, s::Simulation)::Float64
    xm = zeros(Float64, 3)
    xi = zeros(Float64, 3)
    temp = zeros(Float64, 3)
    r = zeros(Float64, 3)
    qf_1 = 0.0
    qf_2 = zeros(Float64, 3)
    for i in 1:N
        for j in 1:12
            V = V_au_au_loop_if(x, V, i, j, xm, xi, temp,r, qf_1, qf_2, s)
        end
    end
    return V
end

@inline function get_V_au_au(s::Simulation)::Float64    #check
    V = 0.0
    V = V_au_au_loop(s.x, V, s)
    return V/4.0
end


text:
module X

using CompEcon

immutable Agent
    ρ::Float64
    α::Float64
    β::Float64
end

_unpack(a::Agent) = (a.ρ, a.α, a.β)

immutable Exog
    A::Float64
    B::Float64
    τ::Float64
    φᵥ::Float64
    vbar::Float64
    p::Matrix{Float64}
    Π::Vector{Float64}
    xp::Matrix{Float64}
    vp::Matrix{Float64}
    gp::Matrix{Float64}
    lgp::Matrix{Float64}
end

function Exog(A::Float64, B::Float64, τ::Float64, φv::Float64, vbar::Float64,
              lgbar::Float64, nϵ1::Int, nϵ2::Int, grid::Matrix{Float64})

    # extract x_t and v_t from grid
    x = grid[:, 2]
    v = grid[:, 3]

    # construct exogenous process
    ϵ, Π = qnwnorm([nϵ1, nϵ2], [0.0, 0.0], eye(2))
    xp = A*x .+ B*sqrt(v).*ϵ[:, 1]'
    vp = (1-φv)*vbar .+ φv*v .+ τ*ϵ[:, 2]'
    gp = exp(lgbar + xp)
    lgp = log(gp)

    # transpose so we can iterate a across states column by column
    xp = xp'
    vp = vp'
    gp = gp'
    lgp = lgp'

    # package
    Exog(A, B, τ, φv, vbar, ϵ, Π, xp, vp, gp, lgp)
end

# ----- #
# Model #
# ----- #
immutable BFZ{BST}
    agent::Agent
    η::Float64
    ν::Float64
    δ::Float64
    grid::Matrix{Float64}
    grid_transpose::Matrix{Float64}  # for iterating state by state over cols
    dim_map::Dict{Symbol,Int}
    basis::Basis{3}
    basis_struc::BasisStructure{BST}
    exog::Exog
end

function BFZ(;ρ::Real=-1, α::Real=-9, β::Real=0.99,
              η::Real=1/3, ν::Float64=0.1, δ::Real=0.025,
              lgbar::Real=0.004, e::Real=1.0, A::Real=0., B::Real=1.,
              vbar::Real=0.015^2, φv::Real=0.95, τ::Real=0.74e-5,
              nk::Int=20, nx::Int=8, nv::Int=8, nϵ1::Int=3, nϵ2::Int=3)

    basis = Basis(SplineParams(collect(linspace(25.0, 50.0, nk)), 0, 3),
                  SplineParams(collect(linspace(-.04, .04, nx)), 0, 1),
                  SplineParams(collect(linspace(1e-4, 4e-4, nv)), 0, 1))

    grid = nodes(basis)[1]

    exog = Exog(A, B, τ, φv, vbar, lgbar, nϵ1, nϵ2, grid)

    basis_struc = BasisStructure(basis, Direct())

    # Create the agent
    agent = Agent(ρ, α, β)

    dim_map = Dict(zip([:k, :x, :v], 1:3))

    # package everything up and return
    BFZ(agent, η, ν, δ, grid, grid', dim_map, basis, basis_struc, exog)
end

_Y(m, k) = (m.η*k.^(m.ν) + (1-m.η)).^(1/m.ν)

function solve_this_state(m::BFZ, i::Int, xv_mat::AbstractMatrix, coefs::Vector,
                          y::Real, j::Real, jk::Real, gp::Vector)
    k, x, v = m.grid_transpose[:, i]
    ρ, α, β = _unpack(m.agent)
    ν, η, δ = m.ν, m.η, m.δ

    N = length(gp)

    # solve for c and k'
    c = (jk/(j^(1-ρ) * (1-β) * (y^(1-ν)*η*k^(ν-1) + 1 - δ))).^(1/(ρ - 1))
    kp = (y + (1-δ)*k - c) ./ gp

    # compute j' before evaluating μ
    kp_basis_mat = CompEcon.evalbase(m.basis.params[1], kp, 0)[1]
    jp = CompEcon.row_kron(xv_mat, kp_basis_mat) * coefs

    μ = dot((gp.*jp).^(α), m.exog.Π).^(1/α)
    jm = ((1-β)*c^ρ + β*μ^ρ)^(1/ρ)

    return jm, c, kp
end

function basis_mat_xv(m::BFZ, x::Vector, v::Vector)
    mat_v = CompEcon.evalbase(m.basis.params[3], v, 0)[1]
    mat_x = CompEcon.evalbase(m.basis.params[2], x, 0)[1]
    CompEcon.row_kron(mat_v, mat_x)
end

function solve_ecm(m::BFZ; ξ::Float64=0.4, tol::Float64=1e-10, maxiter::Int=5000)
    # initialize J
    ρ, α, β = _unpack(m.agent)
    ν, η, δ = m.ν, m.η, m.δ
    M = size(m.grid, 1)
    N = size(m.exog.xp, 1)

    jm_old = (1-β)^(1/ρ)*(_Y(m, m.grid[:, 1]) + (1 - δ)*m.grid[:, 1]) ./ 1000.0
    coefs = funfitxy(m.basis, m.basis_struc, jm_old)[1]

    J_k_basis_struc = BasisStructure(m.basis, Direct(),
                                     nodes(m.basis)[1], [1 0 0])

    xv_mats = [basis_mat_xv(m, m.exog.xp[:, i], m.exog.vp[:, i]) for i=1:M]
    err = 1.0
    it = 0

    jm = similar(jm_old)
    y_all = _Y(m, m.grid[:, 1])

    local out  # declare out local so we can return it after the loop
    while err > tol && it < maxiter
        it += 1
        # compute J at each state
        J = funeval(coefs, m.basis_struc, [0 0 0])
        J_k = funeval(coefs, J_k_basis_struc, [1 0 0])

        # solve at each state in parallel
        ssi(i) = solve_this_state(m, i, xv_mats[i], coefs, y_all[i], J[i],
                                  J_k[i], m.exog.gp[:, i])
        out = map(ssi, 1:M)

        # make sure we have jm::Vector{Float64}
        jm = Float64[i[1] for i in out]

        # now let's get ourselves a new coefficient vector
        coef_hat = funfitxy(m.basis, m.basis_struc, jm)[1]
        coefs = (1-ξ)*coef_hat + ξ*coefs
        err = maxabs(jm - jm_old)
        copy!(jm_old, jm)

        @show it, err

    end
    return coefs, out
end

function main()
    m = BFZ()
    solve_ecm(m)
end

end  # module

max_stars_repo_name: NYUEcon/GrowthModels

max_issues_repo_name: NYUEcon/GrowthModels

max_forks_repo_name: NYUEcon/GrowthModels


text:
# assumes symmetric matrix
function mds(D::AbstractMatrix{T}, p::Integer) where {T<:AbstractFloat}
    @assert D==D'
    N = size(D,1)

    D = D.^2 # MDS works with squared distances

    m = mean(D;dims=1)
    D = Symmetric(-1/2*(D.-m.-m'.+mean(m))) # remove mean

    F = eigen(D, N-p+1:N)
    Σ,X = F.values, F.vectors

    #X = diagm(sqrt(Σ))*X'
    #X[p:-1:1,:] # reorder since eig orders from smallest to largest eigenvalue and we want largest to smallest
    X = X*Diagonal(sqrt.(Σ))
    X[:,p:-1:1] # reorder since eig orders from smallest to largest eigenvalue and we want largest to smallest
end

mds(D::AbstractMatrix, p::Integer) = mds(convert(Matrix{Float64},D),p)


text:
using DrWatson, Test
using BSON, DataFrames, FileIO, JLD2

@testset "Collect Results ($ending)" for ending ∈ ["bson", "jld2"]

###############################################################################
#                        Setup Folder structure                               #
###############################################################################
# %%
cd(@__DIR__)
isdir("testdir") && rm("testdir", recursive=true)
mkdir("testdir")
initialize_project("testdir"; git = false)
quickactivate("testdir")

###############################################################################
#                           Create Dummy Data                                 #
###############################################################################
mkdir(datadir("results"))
cd(datadir("results"))

d = Dict("a" => 1, "b" => "2", "c" => rand(10))
DrWatson.wsave(savename(d)*"."*ending, d)

d = Dict("a" => 3, "b" => "4", "c" => rand(10), "d" => Float64)
DrWatson.wsave(savename(d)*"."*ending, d)

d = Dict("a" => 3, "c" => rand(10), "d" => Float64)
DrWatson.wsave(savename(d)*"."*ending, d)

mkdir("subfolder")
cd("subfolder")

d = Dict("a" => 4., "b" => "twenty" , "d" => Int)
DrWatson.wsave(savename(d)*"."*ending, d)

###############################################################################
#                           Collect Data Into DataFrame                       #
###############################################################################
using Statistics
special_list = [ :lv_mean => data -> mean(data["c"]),
                :lv_var  => data -> var(data["c"])]

black_list = ["c"]

folder = datadir("results")

defaultname = joinpath(dirname(folder), "results_$(basename(folder))."*ending)
isfile(defaultname) && rm(defaultname)
cres = collect_results!(defaultname, folder;
    subfolders = true, special_list=special_list, black_list = black_list)

@test size(cres) == (4, 6)
for n in ("a", "b", "lv_mean")
    @test n ∈ String.(names(cres))
end
@test "c" ∉ names(cres)
@test all(startswith.(cres[!,"path"], projectdir()))

relpathname = joinpath(dirname(folder), "results_relpath_$(basename(folder))."*ending)
cres_relpath = collect_results!(relpathname, folder;
    subfolders = true, special_list=special_list, black_list = black_list,
    rpath = projectdir())
@info all(startswith.(cres[!,"path"], "data"))

###############################################################################
#                           Add another file in a sub sub folder              #
###############################################################################

@test isfile(defaultname)

mkdir("subsubfolder")
cd("subsubfolder")
d = Dict("b" => 35., "d" => Number, "c" => rand(5), "e" => "new_column")
DrWatson.wsave(savename(d)*".bson", d)
DrWatson.wsave(savename(d)*".jld2", d)

cres2 = collect_results!(defaultname, folder;
    subfolders = true, special_list=special_list, black_list = black_list)

@test size(cres2) == (6, 7)
@test all(names(cres) .∈ Ref(names(cres2)))

###############################################################################
#               Test additional syntax for special list                       #
###############################################################################

special_list2 = [ :lv_mean => data -> mean(data["c"]),
                data -> :lv_var  =>  var(data["c"]),
                data -> [:lv_mean2 => mean(data["c"]),
                        :lv_var2  =>  var(data["c"])]]
black_list = ["c"]

folder = datadir("results")
defaultname2 = joinpath(dirname(folder), "results_betterspeciallist."*ending)
isfile(defaultname2) && rm(defaultname2)
cres10 = collect_results!(defaultname2, folder;
    subfolders = true, special_list=special_list2, black_list = black_list)

@test size(cres10) == (6, 9)
for n in ("a", "b", "lv_mean", "lv_var", "lv_mean2", "lv_var2")
    @test n ∈ String.(names(cres10))
end
@test "c" ∉ names(cres10)
###############################################################################
#                           Load and analyze  DataFrame                       #
###############################################################################

df = load(defaultname)["df"]
@test size(df) == size(cres2)
@test sort(names(df)) == sort(names(cres2))

###############################################################################
#                            test empty whitelist                             #
###############################################################################

rm(defaultname)
cres_empty = collect_results!(defaultname, folder;
    subfolders = true, special_list=special_list, white_list=[])

@test dropmissing(cres2[!,[:lv_mean, :lv_var, :path]]) == dropmissing(cres_empty)

###############################################################################
#                           test out-of-place form                            #
###############################################################################
cd(@__DIR__)

cres3 = collect_results(folder;
subfolders = true, special_list=special_list, black_list = black_list)

@test sort(names(cres3)) == sort(names(cres2))
@test size(cres3) == size(cres2)

###############################################################################
#                           test updating feature                             #
###############################################################################

@testset "Test updating feature $(mtime_info)" for mtime_info in ["with mtime", "without initial update", "without mtime", "with corrupt mtime"]
    # Create a temp directory and run the tests, creating files in that folder
    # Julia takes care of removing the folder after the function is done.
    mktempdir(datadir()) do folder
        # Create three data files with slightly different data
        d = Dict("idx" => :keep, "b" => "some_value")
        fname_keep = joinpath(folder, savename(d, ending, ignores = ("b",)))
        DrWatson.wsave(fname_keep, d)

        d = Dict("idx" => :delete, "b" => "some_other_value")
        fname_delete = joinpath(folder, savename(d, ending, ignores = ("b",)))
        DrWatson.wsave(fname_delete, d)

        d = Dict("idx" => :to_modify, "b" => "original_value")
        fname_modify = joinpath(folder, savename(d, ending, ignores = ("b",)))
        DrWatson.wsave(fname_modify, d)

        # Collect our "results"
        if mtime_info == "without initial update"
            # Test this case: https://github.com/JuliaDynamics/DrWatson.jl/pull/286#pullrequestreview-755999610
            cres_before = collect_results!(folder; update = false)
        else
            cres_before = collect_results!(folder; update = true)
        end

        if mtime_info == "without mtime"
            # Leave out the mtime information to simulate old results collection.
            wsave(joinpath(dirname(folder), "results_$(basename(folder)).jld2"), Dict("df" => cres_before))
        elseif mtime_info == "with corrupt mtime"
            # Corrupt mtime information
            wsave(joinpath(dirname(folder), "results_$(basename(folder)).jld2"), Dict("df" => cres_before, "mtime" => Dict{String,Float64}()))
        else
            # Modify one data file
            d = Dict("idx" => :to_modify, "b" => "modified_value")
            DrWatson.wsave(fname_modify, d)

            # Delete another data file
            rm(fname_delete)
        end

        # Collect the "results" again
        if (mtime_info == "without mtime") || (mtime_info == "with corrupt mtime")
            @test_throws DrWatson.InvalidResultsCollection collect_results!(folder; update = true)
        else
            cres_after = collect_results!(folder; update = true)

            # Compare the before and after - they should differ
            @test cres_before[:,[:idx, :b]] != cres_after[:,[:idx, :b]]
            # The unmodified entry should be the same
            @test ((:keep ∈ cres_before.idx) && (:keep ∈ cres_after.idx))
            # The deleted entry should be gone
            @test ((:delete ∈ cres_before.idx) && (:delete ∉ cres_after.idx))
            # The modified entry should differ between before and after
            @test cres_before.b[cres_before.idx .== :to_modify][1] == "original_value"
            @test cres_after.b[cres_after.idx .== :to_modify][1] == "modified_value"
        end
    end
end

###############################################################################
#                              Quickactivate macro                            #
###############################################################################

cd(@__DIR__)
isdir("testdir") && rm("testdir", recursive=true)
initialize_project("testdir"; git = false)
open(joinpath("testdir","testinclude.jl"),"w") do f
    write(f,"@quickactivate\n")
end
include(joinpath("testdir", "testinclude.jl"))
@test Base.active_project() == abspath(joinpath("testdir","Project.toml"))

###############################################################################
#                                 Delete Folders                              #
###############################################################################

cd(@__DIR__)
rm("testdir", recursive=true)

end


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

struct BatchProcessingError <: Exception
    data
    ex
end

"""
    pgenerate([::WorkerPool], f, c...) -> iterator

Apply `f` to each element of `c` in parallel using available workers and tasks.

For multiple collection arguments, apply `f` elementwise.

Results are returned in order as they become available.

Note that `f` must be made available to all worker processes; see
[Code Availability and Loading Packages](@ref code-availability)
for details.
"""
function pgenerate(p::WorkerPool, f, c)
    if length(p) == 0
        return AsyncGenerator(f, c; ntasks=()->nworkers(p))
    end
    batches = batchsplit(c, min_batch_count = length(p) * 3)
    return Iterators.flatten(AsyncGenerator(remote(p, b -> asyncmap(f, b)), batches))
end
pgenerate(p::WorkerPool, f, c1, c...) = pgenerate(p, a->f(a...), zip(c1, c...))
pgenerate(f, c) = pgenerate(default_worker_pool(), f, c)
pgenerate(f, c1, c...) = pgenerate(a->f(a...), zip(c1, c...))

"""
    pmap(f, [::AbstractWorkerPool], c...; distributed=true, batch_size=1, on_error=nothing, retry_delays=[], retry_check=nothing) -> collection

Transform collection `c` by applying `f` to each element using available
workers and tasks.

For multiple collection arguments, apply `f` elementwise.

Note that `f` must be made available to all worker processes; see
[Code Availability and Loading Packages](@ref code-availability) for details.

If a worker pool is not specified, all available workers, i.e., the default worker pool
is used.

By default, `pmap` distributes the computation over all specified workers. To use only the
local process and distribute over tasks, specify `distributed=false`.
This is equivalent to using [`asyncmap`](@ref). For example,
`pmap(f, c; distributed=false)` is equivalent to `asyncmap(f,c; ntasks=()->nworkers())`

`pmap` can also use a mix of processes and tasks via the `batch_size` argument. For batch sizes
greater than 1, the collection is processed in multiple batches, each of length `batch_size` or less.
A batch is sent as a single request to a free worker, where a local [`asyncmap`](@ref) processes
elements from the batch using multiple concurrent tasks.

Any error stops `pmap` from processing the remainder of the collection. To override this behavior
you can specify an error handling function via argument `on_error` which takes in a single argument, i.e.,
the exception. The function can stop the processing by rethrowing the error, or, to continue, return any value
which is then returned inline with the results to the caller.

Consider the following two examples. The first one returns the exception object inline,
the second a 0 in place of any exception:
```julia-repl
julia> pmap(x->iseven(x) ? error("foo") : x, 1:4; on_error=identity)
4-element Array{Any,1}:
 1
  ErrorException("foo")
 3
  ErrorException("foo")

julia> pmap(x->iseven(x) ? error("foo") : x, 1:4; on_error=ex->0)
4-element Array{Int64,1}:
 1
 0
 3
 0
```

Errors can also be handled by retrying failed computations. Keyword arguments `retry_delays` and
`retry_check` are passed through to [`retry`](@ref) as keyword arguments `delays` and `check`
respectively. If batching is specified, and an entire batch fails, all items in
the batch are retried.

Note that if both `on_error` and `retry_delays` are specified, the `on_error` hook is called
before retrying. If `on_error` does not throw (or rethrow) an exception, the element will not
be retried.

Example: On errors, retry `f` on an element a maximum of 3 times without any delay between retries.
```julia
pmap(f, c; retry_delays = zeros(3))
```

Example: Retry `f` only if the exception is not of type [`InexactError`](@ref), with exponentially increasing
delays up to 3 times. Return a `NaN` in place for all `InexactError` occurrences.
```julia
pmap(f, c; on_error = e->(isa(e, InexactError) ? NaN : rethrow()), retry_delays = ExponentialBackOff(n = 3))
```
"""
function pmap(f, p::AbstractWorkerPool, c; distributed=true, batch_size=1, on_error=nothing,
                                           retry_delays=[], retry_check=nothing)
    f_orig = f
    # Don't do remote calls if there are no workers.
    if (length(p) == 0) || (length(p) == 1 && fetch(p.channel) == myid())
        distributed = false
    end

    # Don't do batching if not doing remote calls.
    if !distributed
        batch_size = 1
    end

    # If not batching, do simple remote call.
    if batch_size == 1
        if on_error !== nothing
            f = wrap_on_error(f, on_error)
        end

        if distributed
            f = remote(p, f)
        end

        if length(retry_delays) > 0
            f = wrap_retry(f, retry_delays, retry_check)
        end

        return asyncmap(f, c; ntasks=()->nworkers(p))
    else
        # During batch processing, We need to ensure that if on_error is set, it is called
        # for each element in error, and that we return as many elements as the original list.
        # retry, if set, has to be called element wise and we will do a best-effort
        # to ensure that we do not call mapped function on the same element more than length(retry_delays).
        # This guarantee is not possible in case of worker death / network errors, wherein
        # we will retry the entire batch on a new worker.

        handle_errors = ((on_error !== nothing) || (length(retry_delays) > 0))

        # Unlike the non-batch case, in batch mode, we trap all errors and the on_error hook (if present)
        # is processed later in non-batch mode.
        if handle_errors
            f = wrap_on_error(f, (x,e)->BatchProcessingError(x,e); capture_data=true)
        end

        f = wrap_batch(f, p, handle_errors)
        results = asyncmap(f, c; ntasks=()->nworkers(p), batch_size=batch_size)

        # process errors if any.
        if handle_errors
            process_batch_errors!(p, f_orig, results, on_error, retry_delays, retry_check)
        end

        return results
    end
end

pmap(f, p::AbstractWorkerPool, c1, c...; kwargs...) = pmap(a->f(a...), p, zip(c1, c...); kwargs...)
pmap(f, c; kwargs...) = pmap(f, default_worker_pool(), c; kwargs...)
pmap(f, c1, c...; kwargs...) = pmap(a->f(a...), zip(c1, c...); kwargs...)

function wrap_on_error(f, on_error; capture_data=false)
    return x -> begin
        try
            f(x)
        catch e
            if capture_data
                on_error(x, e)
            else
                on_error(e)
            end
        end
    end
end

function wrap_retry(f, retry_delays, retry_check)
    retry(delays=retry_delays, check=retry_check) do x
        try
            f(x)
        catch e
            rethrow(extract_exception(e))
        end
    end
end

function wrap_batch(f, p, handle_errors)
    f = asyncmap_batch(f)
    return batch -> begin
        try
            remotecall_fetch(f, p, batch)
        catch e
            if handle_errors
                return Any[BatchProcessingError(b, e) for b in batch]
            else
                rethrow()
            end
        end
    end
end

asyncmap_batch(f) = batch -> asyncmap(x->f(x...), batch)
extract_exception(e) = isa(e, RemoteException) ? e.captured.ex : e

function process_batch_errors!(p, f, results, on_error, retry_delays, retry_check)
    # Handle all the ones in error in another pmap, with batch size set to 1
    reprocess = []
    for (idx, v) in enumerate(results)
        if isa(v, BatchProcessingError)
            push!(reprocess, (idx,v))
        end
    end

    if length(reprocess) > 0
        errors = [x[2] for x in reprocess]
        exceptions = [x.ex for x in errors]
        state = iterate(retry_delays)
        state !== nothing && (state = state[2])
        if (length(retry_delays) > 0) &&
                (retry_check === nothing || all([retry_check(state,ex)[2] for ex in exceptions]))
            # BatchProcessingError.data is a tuple of original args
            error_processed = pmap(x->f(x...), p, [x.data for x in errors];
                    on_error = on_error, retry_delays = collect(retry_delays)[2:end], retry_check = retry_check)
        elseif on_error !== nothing
            error_processed = map(on_error, exceptions)
        else
            throw(CompositeException(exceptions))
        end

        for (idx, v) in enumerate(error_processed)
            results[reprocess[idx][1]] = v
        end
    end
    nothing
end

"""
    head_and_tail(c, n) -> head, tail

Return `head`: the first `n` elements of `c`;
and `tail`: an iterator over the remaining elements.

```jldoctest
julia> b, c = Distributed.head_and_tail(1:10, 3)
([1, 2, 3], Base.Iterators.Rest{UnitRange{Int64},Int64}(1:10, 3))

julia> collect(c)
7-element Vector{Int64}:
  4
  5
  6
  7
  8
  9
 10
```
"""
function head_and_tail(c, n)
    head = Vector{eltype(c)}(undef, n)
    n == 0 && return (head, c)
    i = 1
    y = iterate(c)
    y === nothing && return (resize!(head, 0), ())
    head[i] = y[1]
    while i < n
        y = iterate(c, y[2])
        y === nothing && return (resize!(head, i), ())
        i += 1
        head[i] = y[1]
    end
    return head, Iterators.rest(c, y[2])
end

"""
    batchsplit(c; min_batch_count=1, max_batch_size=100) -> iterator

Split a collection into at least `min_batch_count` batches.

Equivalent to `partition(c, max_batch_size)` when `length(c) >> max_batch_size`.
"""
function batchsplit(c; min_batch_count=1, max_batch_size=100)
    if min_batch_count < 1
        throw(ArgumentError("min_batch_count must be ≥ 1, got $min_batch_count"))
    end

    if max_batch_size < 1
        throw(ArgumentError("max_batch_size must be ≥ 1, got $max_batch_size"))
    end

    # Split collection into batches, then peek at the first few batches
    batches = Iterators.partition(c, max_batch_size)
    head, tail = head_and_tail(batches, min_batch_count)

    # If there are not enough batches, use a smaller batch size
    if length(head) < min_batch_count
        batch_size = max(1, div(sum(length, head), min_batch_count))
        return Iterators.partition(collect(Iterators.flatten(head)), batch_size)
    end

    return Iterators.flatten((head, tail))
end


text:
using ProgressMeter, InfoZIP
using HTTP.IOExtras
using Base.Filesystem: IOError

function post_with_err_handling(url, body, headers; timeout=0)
  try
    resp = HTTP.post(url, headers=headers, body=JSON.json(body), readtimeout=timeout)
    return JSON.parse(String(resp.body))
  catch e
    return Dict()
  end
end

function parse_datetime_string(datetime_string; datetime_format="yyyy-mm-ddTHH:MM:SS")
  return isnothing(datetime_string) ? datetime_string : DateTime(datetime_string[1:end-1], datetime_format)
end

function parse_float_string(float_string)
  return isnothing(float_string) ? float_string : tryparse(Float64, float_string)
end

function ensure_directory_exists(path)
  try
    mkdir(path)
  catch e
    if !isa(e, IOError)
      throw(e)
    end
  end
end

function download_file(url, dest_path; show_progress_bar=true)

  HTTP.open("GET", url) do http
    r = startread(http)
    l = parse(Int, HTTP.header(r, "Content-Length"))

    if show_progress_bar
      prog = Progress(l, dt=0.1, desc="Downloading:", barglyphs=BarGlyphs("[=> ]"), barlen=25, color=:white)
    end

    io = open(dest_path, "w")
    n_bytes = 0
    while !eof(http)
      bytes = readavailable(http)
      n_bytes += length(bytes)
      write(io, bytes)
      if show_progress_bar
        ProgressMeter.update!(prog, n_bytes)
      end
    end
    close(io)
  end

end

function unzip_file(src_path::String, dest_path::String, filename::String)
  @info "unzipping file..."

  unzip_path = joinpath(dest_path, filename)
  ensure_directory_exists(unzip_path)

  InfoZIP.unzip(src_path, unzip_path)

end


text:
using Distributions, StatsPlots, Random

# Set a random seed.
Random.seed!(3)

# Construct 30 data points for each cluster.
N = 30

# Parameters for each cluster, we assume that each cluster is Gaussian distributed in the example.
μs = [-3.5, 0.0]

# Construct the data points.
x = mapreduce(c -> rand(MvNormal([μs[c], μs[c]], 1.), N), hcat, 1:2)

# Visualization.
scatter(x[1,:], x[2,:], legend = false, title = "Synthetic Dataset")

using Turing, MCMCChains

# Turn off the progress monitor.
Turing.setprogress!(false);

@model GaussianMixtureModel(x) = begin

    D, N = size(x)

    # Draw the parameters for cluster 1.
    μ1 ~ Normal()

    # Draw the parameters for cluster 2.
    μ2 ~ Normal()

    μ = [μ1, μ2]

    # Uncomment the following lines to draw the weights for the K clusters
    # from a Dirichlet distribution.

    # α = 1.0
    # w ~ Dirichlet(2, α)

    # Comment out this line if you instead want to draw the weights.
    w = [0.5, 0.5]

    # Draw assignments for each datum and generate it from a multivariate normal.
    k = Vector{Int}(undef, N)
    for i in 1:N
        k[i] ~ Categorical(w)
        x[:,i] ~ MvNormal([μ[k[i]], μ[k[i]]], 1.)
    end
    return k
end

gmm_model = GaussianMixtureModel(x);

gmm_sampler = Gibbs(PG(100, :k), HMC(0.05, 10, :μ1, :μ2))
tchain = mapreduce(c -> sample(gmm_model, gmm_sampler, 100), chainscat, 1:3);

ids = findall(map(name -> occursin("μ", string(name)), names(tchain)));
p=plot(tchain[:, ids, :], legend=true, labels = ["Mu 1" "Mu 2"], colordim=:parameter)

tchain = tchain[:, :, 1];

# Helper function used for visualizing the density region.
function predict(x, y, w, μ)
    # Use log-sum-exp trick for numeric stability.
    return Turing.logaddexp(
        log(w[1]) + logpdf(MvNormal([μ[1], μ[1]], 1.), [x, y]),
        log(w[2]) + logpdf(MvNormal([μ[2], μ[2]], 1.), [x, y])
    )
end

contour(range(-5, stop = 3), range(-6, stop = 2),
    (x, y) -> predict(x, y, [0.5, 0.5], [mean(tchain[:μ1]), mean(tchain[:μ2])])
)
scatter!(x[1,:], x[2,:], legend = false, title = "Synthetic Dataset")

# TODO: is there a better way than this icky `.nt.mean` stuff?
assignments = mean(MCMCChains.group(tchain, :k)).nt.mean
scatter(x[1,:], x[2,:],
    legend = false,
    title = "Assignments on Synthetic Dataset",
    zcolor = assignments)

if isdefined(Main, :TuringTutorials)
    Main.TuringTutorials.tutorial_footer(WEAVE_ARGS[:folder], WEAVE_ARGS[:file])
end

max_stars_repo_name: leachim/TuringTutorials

max_issues_repo_name: leachim/TuringTutorials

max_forks_repo_name: leachim/TuringTutorials


text:
# This example predicts the temporal evolution of a one-dimensional field
# U, along with a time vector T,
# which has to be represented as vectors of vectors.
# Where the field comes from does not matter, but to make the example
# runnable we load one of the test systems of `TimeseriesPrediction`.
#
# In this example we use the solution of Kuramoto Sivashinsky equation.
#
# Importantly, the results are compared with the "real" evolution of the
# system.
#
# In the plots, the x axis is space and y axis is time.

# ### Produce field U (Kuramoto Sivashinsky)
using PyPlot
using TimeseriesPrediction

testdir = dirname(dirname(pathof(TimeseriesPrediction)))*"/test"
@assert isdir(testdir)
include(testdir*"/ks_solver.jl")

Ntrain = 10000
p = 100
N = Ntrain + p

U, T = KuramotoSivashinsky(64, 22, N÷4, 0.25)
summary(U)

# ### Temporal prediction of field U
Q = length(U[1]) # spatial length
pool = U[1:Ntrain]
test = U[Ntrain:N]

γ = 10
τ = 1
B = 10
k = 1
ntype = FixedMassNeighborhood(4)
method = AverageLocalModel()

em = cubic_shell_embedding(pool, γ,τ,B,k,PeriodicBoundary())
pcaem= PCAEmbedding(pool,em)

@time pred = temporalprediction(pool,pcaem, p;ntype=ntype, method=method, progress = false)

err = [abs.(test[i]-pred[i]) for i=1:p+1]
println("Maximum error: ", maximum(maximum(e) for e in err))

# ### Plot the result

# Deduce field extremal values
vmax = max(maximum(maximum(s) for s in test),
           maximum(maximum(s) for s in pred))
vmin = min(minimum(minimum(s) for s in test),
           minimum(minimum(s) for s in pred))

# Transform data for imshow
ptest = cat(test..., dims = 2)
ppred = cat(pred..., dims = 2)
perr = cat(err..., dims = 2)

# plot plot plot
fig = figure(figsize=(8,8))
ax1 = subplot2grid((3,1), (0,0))
ax2 = subplot2grid((3,1), (1,0))
ax3 = subplot2grid((3,1), (2,0));

im1 = ax1[:imshow](ppred, cmap="viridis", vmin = vmin, vmax = vmax,
aspect = "auto", extent = (T[Ntrain], T[N], 1, Q))
im2 = ax2[:imshow](ptest, cmap="viridis", vmin = vmin, vmax = vmax,
aspect = "auto", extent = (T[Ntrain], T[N], 1, Q))
im3 = ax3[:imshow](perr, cmap="inferno", vmin = 0, vmax = vmax-vmin,
aspect = "auto", extent = (T[Ntrain], T[N], 1, Q))

for (j, (im, ax)) in enumerate(zip([im1,im2,im3], [ax1,ax2,ax3]))
    colorbar(im, ax = ax, fraction=0.04, pad=0.01)# format="%.1f")
    if j < 3
        ax[:set_xticklabels]([])
    end
end
ax1[:set_title]("Prediction")
ax2[:set_title]("Real evolution")
ax3[:set_title]("Absolute error")

ax2[:set_ylabel]("space")
ax3[:set_xlabel]("time")
tight_layout(w_pad=0.1, h_pad=0.00001)
#md savefig("ksprediction.png"); nothing # hide
#md # ![ksprediction](ksprediction.png)


text:
# Michelle Medina
# August 3, 2020

export derivBasis, derivEvaluate

function derivBasis(l::Int, j::Int, x::T)::T where {T}
	@assert 0 <= l
	@assert 0 <= j <= (2^l)
	h = 2.0^(-l)
	eps = 1e-10
	if 0 <= x < eps
		if j == 0
			return -1/h
		elseif j == 1
			return 1/h
		else
			return 0
		end
	elseif 1 - eps < x <= 1
		if j == 2^l
			return 1/h
		elseif j == 2^l - 1
			return -1/h
		else
			return 0
		end
	elseif eps <= x < (j-1)*h - eps
		return 0
	elseif (j-1)*h - eps <= x < (j-1)*h + eps
		return 1/(2*h)
	elseif (j-1)*h + eps <= x < j*h - eps
		return 1/h
	elseif j*h - eps <= x < j*h + eps
		return 0
	elseif j*h + eps <= x < (j+1)*h - eps
		return -1/h
	elseif (j+1)*h - eps <= x < (j+1)*h + eps
		return -1/(2*h)
	else (j+1)*h + eps <= x <= 1 - eps
		return 0
	end
	return value
end

function derivEvaluate(H::HierarchicalBasis{T}, x::T)::T where {T}
	value = T(0)
	for l in 0:maxlevel(H)
		for j in 0:2^l
			value = value + H[l, j]*derivBasis(l, j, x)
		end
	end
	return value
end


text:
using DrWatson
@quickactivate
using ArgParse
using GenerativeAD
import StatsBase: fit!, predict
using StatsBase
using BSON
# because of vae and 2stage
using DataFrames
using CSV
using ValueHistories
using Flux
using ConditionalDists
using GenerativeModels
import GenerativeModels: VAE
using Distributions
using DistributionsAD
using Distances
using Statistics

s = ArgParseSettings()
@add_arg_table! s begin
   "max_seed"
		required = true
		arg_type = Int
		help = "seed"
	"dataset"
		required = true
		arg_type = String
		help = "dataset"
	"anomaly_class"
		arg_type = Int
		default = 1
		help = "anomaly classes"
	"method"
		arg_type = String
		default = "leave-one-out"
		help = "method for data creation -> \"leave-one-out\" or \"leave-one-in\" "
    "contamination"
    	arg_type = Float64
    	help = "contamination rate of training data"
    	default = 0.0
end
parsed_args = parse_args(ARGS, s)
@unpack dataset, anomaly_class, max_seed, method, contamination = parsed_args

#######################################################################################
################ THIS PART IS TO BE PROVIDED FOR EACH MODEL SEPARATELY ################

tab_name = "vae_LOSS_images_mnistc"
mi = 1
sp = split(tab_name, "_")
enc = sp[1]
criterion = lowercase(sp[2])
modelname = "$(enc)_ocsvm"

function median_l2_dist(X)
	dists = pairwise(Euclidean(), X)
	# take only the upper diagonal
	ds = []
	for i in 1:size(dists,1)
		for j in (i+1):size(dists,1)
			push!(ds, dists[i,j])
		end
	end
	ml=median(ds)
end
function set_params(data)
	D, N = size(data[1][1])
	gamma = 1/median_l2_dist(data[1][1])
	return (gamma=gamma, kernel="rbf", nu=0.5)
end

function joint_fit(models, data_splits)
	info = []
	for (model, data) in zip(models, data_splits)
		push!(info, fit!(model, data))
	end
	return info
end

function joint_prediction(models, data)
	joint_pred = Array{Float32}(undef, length(models), size(data,2))
	for (i,model) in enumerate(models)
		joint_pred[i,:] = predict(model, data)
	end
	return vec(mean(joint_pred, dims=1))
end

function StatsBase.fit(data, parameters, n_models, aux_info)
	# construct model - constructor should only accept kwargs
	models = [GenerativeAD.Models.OCSVM(;parameters...) for _ = 1:n_models]

	# sumbsample and fit train data
	tr_data = data[1][1]
	M,N = size(tr_data)
	split_size = floor(N/n_models)
	data_splits = [tr_data[:,Int(split_size*i+1):Int(split_size*(i+1))] for i = 0:n_models-1]

	try
		global info, fit_t, _, _, _ = @timed joint_fit(models, data_splits)
 	catch e
		# return an empty array if fit fails so nothing is computed
		return (fit_t = NaN,), []
	end

	# construct return information - put e.g. the model structure here for generative models
	training_info = (
		fit_t = fit_t,
		model = nothing
		)

	# now return the different scoring functions
	training_info, [(x->joint_prediction(models, x),  merge(parameters, aux_info))]
end

####################################################################
################ THIS PART IS COMMON FOR ALL MODELS ################
# set a maximum for parameter sampling retries
try_counter = 0
max_tries = 10*max_seed
cont_string = (contamination == 0.0) ? "" : "_contamination-$contamination"
while try_counter < max_tries
	for seed in 1:max_seed
		i = anomaly_class
		savepath = datadir("experiments/images_mnistc_clean_val_default$cont_string/$(modelname)/$(dataset)/ac=$(i)/seed=$(seed)")
		aux_info = (model_index=mi, criterion=criterion)

		global data = GenerativeAD.load_data(dataset, seed=seed, anomaly_class_ind=i, method=method, contamination=contamination)
		not_loaded = true
		while not_loaded
			try
				global data, encoding_name, encoder_params = GenerativeAD.Models.load_encoding(tab_name, data, i, dataset=dataset, seed=seed, model_index=mi)
				not_loaded = false
			catch e
				@info "model index $mi not working, trying the next one"
				global mi += 1
			end
		end
		parameters = set_params(data)

		# here, check if a model with the same parameters was already tested
		@info "Trying to fit $modelname on $dataset with parameters $(parameters)..."
		if GenerativeAD.check_params(savepath, merge(parameters, aux_info))
			training_info, results = fit(data, parameters, 10, aux_info)
			# here define what additional info should be saved together with parameters, scores, labels and predict times
			save_entries = merge(training_info, (modelname = modelname,
												 seed = seed,
												 dataset = dataset,
												 anomaly_class = i,
												 encoder=encoding_name,
												 encoder_params=encoder_params,
												 model_index=mi,
												 criterion=criterion,
												 contamination=contamination))
			# now loop over all anomaly score funs
			for result in results
				GenerativeAD.experiment(result..., data, savepath; save_entries...)
			end
			global try_counter = max_tries + 1
		else
			@info "Model already present, sampling new hyperparameters..."
			global try_counter += 1
		end
	end
end
(try_counter == max_tries) ? (@info "Reached $(max_tries) tries, giving up.") : nothing