text:
using JLD, Knet
MAJOR=1 # 1:column-major 2:row-major

typealias Word String           # represent words with strings instead of finite vocab, we will process chars
typealias PosTag UInt8          # 17 universal part-of-speech tags
typealias DepRel UInt8          # 37 universal dependency relations
typealias Position Int16        # sentence position
typealias Cost Position         # [0:nword]
typealias Move Int              # [1:nmove]
include("../src/parser.jl")

;                               # CONLLU FORMAT
immutable Sentence		# 1. ID: Word index, integer starting at 1 for each new sentence; may be a range for multiword tokens; may be a decimal number for empty nodes.
    word::Vector{Word}          # 2. FORM: Word form or punctuation symbol.
    #stem::Vector{Stem}         # 3. LEMMA: Lemma or stem of word form.
    postag::Vector{PosTag}      # 4. UPOSTAG: Universal part-of-speech tag.
    #xpostag::Vector{Xpostag}   # 5. XPOSTAG: Language-specific part-of-speech tag; underscore if not available.
    #feats::Vector{Feats}       # 6. FEATS: List of morphological features from the universal feature inventory or from a defined language-specific extension; underscore if not available.
    head::Vector{Position}      # 7. HEAD: Head of the current word, which is either a value of ID or zero (0).
    deprel::Vector{DepRel}      # 8. DEPREL: Universal dependency relation to the HEAD (root iff HEAD = 0) or a defined language-specific subtype of one.
    #deps::Vector{Deps}         # 9. DEPS: Enhanced dependency graph in the form of a list of head-deprel pairs.
    #misc::Vector{Misc}         # 10. MISC: Any other annotation.
    #vocab::Vocab               # Common repository of symbols for upostag, deprel etc.
end
Sentence()=Sentence([],[],[],[])
Base.length(s::Sentence)=length(s.word)

immutable Vocab
    cdict::Dict{Char,Int}       # character dictionary (input)
    idict::Dict{String,Int}     # word dictionary (input, filled in by maptoint)
    odict::Dict{String,Int}     # word dictionary (output)
    sosword::String             # start-of-sentence word (input)
    eosword::String             # end-of-sentence word (input)
    unkword::String             # unknown word (output, input does not have unk)
    sowchar::Char               # start-of-word char
    eowchar::Char               # end-of-word char
    unkchar::Char               # unknown char
    postags::Dict{String,PosTag}
    deprels::Dict{String,DepRel}
end

type BeamState
    parent::BeamState   # previous state
    move::Move          # move executed to get to this state
    score::Float64      # cumulative score (logp) TODO: should be unnormalized logp!
    cost::Cost          # cumulative cost (gold arcs that have become impossible)
    fidx::Int           # column in feature matrix representing current parser state # TODO: what is this?
    parser::Parser 	# current parser state
    BeamState()=new()
    BeamState(p::Parser)=new(NullBeamState,0,0,0,0,p)
    BeamState(b::BeamState,m::Integer,s::Number,c::Integer)=
        new(b,convert(Move,m),convert(Float64,s),convert(Cost,c))
end

const NullBeamState = BeamState()
Base.isless(a::BeamState,b::BeamState)=(a.score < b.score)

type Beam
    beam::Vector{BeamState}
    cand::Vector{BeamState}
    sent::Sentence
    stop::Bool                  # TODO: what is this?
    function Beam{T<:Parser}(pt::Type{T}, s::Sentence)
        p = pt(s)
        b = BeamState(p)
        new(BeamState[b], BeamState[], s, !anyvalidmoves(p))
    end
end

typealias Batch Vector{Beam}

# parse a minibatch of sentences using beam search, global normalization, early updates and report loss.
function parseloss(model, sentences, vocab, parsertype, beamsize)
    # global wembed,forw,back,cdata,cmask,wdata,wmask
    words,sents,wordlen,sentlen = maptoint(sentences, vocab)
    # Get word embeddings
    sow,eow = vocab.cdict[vocab.sowchar],vocab.cdict[vocab.eowchar]
    cdata,cmask = tokenbatch(words,wordlen,sow,eow) # cdata/cmask[T+2][V] where T: longest word (140), V: input word vocab (19674)
    wembed = charlstm(model,cdata,cmask)            # wembed[C,V] where V: input word vocab, C: charlstm hidden (350)
    # Get context embeddings
    sos,eos = vocab.idict[vocab.sosword],vocab.idict[vocab.eosword]
    wdata,wmask = tokenbatch(sents,sentlen,sos,eos) # wdata/wmask[T+2][B] where T: longest sentence (159), B: batch size (12543)
    forw,back = wordlstm(model,wdata,wmask,wembed)  # forw/back[T][W,B] where T: longest sentence, B: batch size, W: wordlstm hidden (300)
    nothing
end

# function loss(model, sentences, parsertype, beamsize)
function lmloss(model, sentences, vocab::Vocab; result=nothing)
    # map words and chars to Ints
    # words and sents contain Int (char/word id) arrays without sos/eos tokens
    # TODO: optionally construct cdict here as well
    words,sents,wordlen,sentlen = maptoint(sentences, vocab)

    # Find word vectors for a batch of sentences using character lstm or cnn model
    # note that tokenbatch adds sos/eos tokens and padding to each sequence
    sow,eow = vocab.cdict[vocab.sowchar],vocab.cdict[vocab.eowchar]
    cdata,cmask = tokenbatch(words,wordlen,sow,eow)
    wembed = charlstm(model,cdata,cmask)

    # Find context vectors using word blstm
    sos,eos = vocab.idict[vocab.sosword],vocab.idict[vocab.eosword]
    wdata,wmask = tokenbatch(sents,sentlen,sos,eos)
    forw,back = wordlstm(model,wdata,wmask,wembed)

    # Test predictions
    unk = vocab.odict[vocab.unkword]
    odata,omask = goldbatch(sentences,sentlen,vocab.odict,unk)
    total = lmloss1(model,odata,omask,forw,back; result=result)
    return -total/length(sentences)
end

function lmloss1(model,data,mask,forw,back; result=nothing)
    # data[t][b]::Int gives the correct word at sentence b, position t
    T = length(data); if !(T == length(forw) == length(back) == length(mask)); error(); end
    B = length(data[1])
    weight,bias = wsoft(model),bsoft(model)
    prd(t) = (if MAJOR==1; weight * vcat(forw[t],back[t]) .+ bias; else; hcat(forw[t],back[t]) * weight .+ bias; end)
    idx(t,b,n) = (if MAJOR==1; data[t][b] + (b-1)*n; else; b + (data[t][b]-1)*n; end)
    total = count = 0
    @inbounds for t=1:T
        ypred = prd(t)
        nrows,ncols = size(ypred)
        index = Int[]
        for b=1:B
            if mask[t][b]==1
                push!(index, idx(t,b,nrows))
            end
        end
        o1 = logp(ypred,MAJOR)
        o2 = o1[index]
        total += sum(o2)
        count += length(o2)
    end
    if result != nothing; result[1]+=AutoGrad.getval(total); result[2]+=count; end
    return total
end

# Find all unique words, assign an id to each, and lookup their characters in cdict
# TODO: construct/add-to cdict here as well?
# in vocab: reads sosword, eosword, unkchar, cdict; writes idict
function maptoint(sentences, v::Vocab)
    wdict = empty!(v.idict)
    cdict = v.cdict
    unkcid = cdict[v.unkchar]
    words = Vector{Int}[]
    sents = Vector{Int}[]
    wordlen = 0
    sentlen = 0
    @inbounds for w in (v.sosword,v.eosword)
        wid = get!(wdict, w, 1+length(wdict))
        word = Array(Int, length(w))
        wordi = 0
        for c in w
            word[wordi+=1] = get(cdict, c, unkcid)
        end
        if wordi != length(w); error(); end
        if wordi > wordlen; wordlen = wordi; end
        push!(words, word)
    end
    @inbounds for s in sentences
        sent = Array(Int, length(s.word))
        senti = 0
        for w in s.word
            ndict = length(wdict)
            wid = get!(wdict, w, 1+ndict)
            sent[senti+=1] = wid
            if wid == 1+ndict
                word = Array(Int, length(w))
                wordi = 0
                for c in w
                    word[wordi+=1] = get(cdict, c, unkcid)
                end
                if wordi != length(w); error(); end
                if wordi > wordlen; wordlen = wordi; end
                push!(words, word)
            end
        end
        if senti != length(s.word); error(); end
        if senti > sentlen; sentlen = senti; end
        push!(sents, sent)
    end
    if length(wdict) != length(words); error("wdict=$(length(wdict)) words=$(length(words))"); end
    return words,sents,wordlen,sentlen
end

# Create token batches, adding start/end tokens and masks, pad at the beginning
function tokenbatch(sequences,maxlen,sos,eos,pad=eos)
    B = length(sequences)
    T = maxlen + 2
    data = [ Array(Int,B) for t in 1:T ]
    mask = [ Array(Float32,B) for t in 1:T ]
    @inbounds for t in 1:T
        for b in 1:B
            N = length(sequences[b])
            n = t - T + N + 1
            if n < 0
                mask[t][b] = 0
                data[t][b] = pad
            else
                mask[t][b] = 1
                if n == 0
                    data[t][b] = sos
                elseif n <= N
                    data[t][b] = sequences[b][n]
                elseif n == N+1
                    data[t][b] = eos
                else
                    error()
                end
            end
        end
    end
    return data,mask
end

function goldbatch(sentences, maxlen, wdict, unkwid, pad=unkwid)
    B = length(sentences)
    T = maxlen # no need for sos/eos for gold
    data = [ Array(Int,B) for t in 1:T ]
    mask = [ Array(Float32,B) for t in 1:T ]
    @inbounds for t in 1:T
        for b in 1:B
            N = length(sentences[b])
            n = t - T + N
            if n <= 0
                mask[t][b] = 0
                data[t][b] = pad
            else
                mask[t][b] = 1
                data[t][b] = get(wdict, sentences[b].word[n], unkwid)
            end
        end
    end
    return data,mask
end

# Run charid arrays through the LSTM, collect last hidden state as word embedding
function charlstm(model,data,mask)
    weight,bias,embeddings = wchar(model),bchar(model),echar(model)
    T = length(data)
    B = length(data[1])
    H = div(length(bias),4)
    cembed(t)=(if MAJOR==1; embeddings[:,data[t]]; else; embeddings[data[t],:]; end)
    czero=(if MAJOR==1; fill!(similar(bias,H,B), 0); else; fill!(similar(bias,B,H), 0); end)
    hidden = cell = czero       # TODO: cache this
    mask = map(KnetArray, mask) # TODO: dont hardcode atype
    @inbounds for t in 1:T
        (hidden,cell) = lstm(weight,bias,hidden,cell,cembed(t);mask=mask[t])
    end
    return hidden
end

function wordlstm(model,data,mask,embeddings) # col-major
    weight,bias = wforw(model),bforw(model)
    T = length(data)
    B = length(data[1])
    H = div(length(bias),4)
    wembed(t)=(if MAJOR==1; embeddings[:,data[t]]; else; embeddings[data[t],:]; end)
    wzero=(if MAJOR==1; fill!(similar(bias,H,B), 0); else; fill!(similar(bias,B,H), 0); end)
    hidden = cell = wzero
    mask = map(KnetArray, mask)           # TODO: dont hardcode atype
    forw = Array(Any,T-2)       # exclude sos/eos
    @inbounds for t in 1:T-2
        (hidden,cell) = lstm(weight,bias,hidden,cell,wembed(t); mask=mask[t])
        forw[t] = hidden
    end
    weight,bias = wback(model),bback(model)
    if H != div(length(bias),4); error(); end
    hidden = cell = wzero
    back = Array(Any,T-2)
    @inbounds for t in T:-1:3
        (hidden,cell) = lstm(weight,bias,hidden,cell,wembed(t); mask=mask[t])
        back[t-2] = hidden
    end
    return forw,back
end

function lstm(weight,bias,hidden,cell,input; mask=nothing)
    if MAJOR==1
        gates   = weight * vcat(input,hidden) .+ bias
        H       = size(hidden,1)
        forget  = sigm(gates[1:H,:])
        ingate  = sigm(gates[1+H:2H,:])
        outgate = sigm(gates[1+2H:3H,:])
        change  = tanh(gates[1+3H:4H,:])
        if mask!=nothing; mask=reshape(mask,1,length(mask)); end
    else
        gates   = hcat(input,hidden) * weight .+ bias
        H       = size(hidden,2)
        forget  = sigm(gates[:,1:H])
        ingate  = sigm(gates[:,1+H:2H])
        outgate = sigm(gates[:,1+2H:3H])
        change  = tanh(gates[:,1+3H:end])
    end
    cell    = cell .* forget + ingate .* change
    hidden  = outgate .* tanh(cell)
    if mask != nothing
        hidden = hidden .* mask
        cell = cell .* mask
    end
    return (hidden,cell)
end

function loadvocab()
    a = load("english_vocabs.jld")
    global _vocab = Vocab(a["char_vocab"],Dict{String,Int}(),a["word_vocab"],
                          "<s>","</s>","<unk>",'↥','Ϟ','⋮',
                          UPOSTAG, UDEPREL)
end

function loadcorpus(v::Vocab)
    global _corpus = Any[]
    s = Sentence()
    for line in eachline("en-ud-train.conllu")
        if line == "\n"
            push!(_corpus, s)
            s = Sentence()
        elseif (m = match(r"^\d+\t(.+?)\t.+?\t(.+?)\t.+?\t.+?\t(.+?)\t(.+?)\t", line)) != nothing
            #                id   word   lem  upos   xpos feat head   deprel
            push!(s.word, m.captures[1])
            push!(s.postag, v.postags[m.captures[2]])
            push!(s.head, parse(Position,m.captures[3]))
            push!(s.deprel, v.deprels[m.captures[4]])
        end
    end
end

function loadmodel()
    global _model = Dict()
    m = load("english_ch12kmodel_2.jld")["model"]
    if MAJOR==1
        for k in (:cembed,); _model[k] = KnetArray(m[k]'); end
        for k in (:forw,:soft,:back,:char); _model[k] = map(KnetArray, m[k]'); end
    else
        for k in (:cembed,); _model[k] = KnetArray(m[k]); end
        for k in (:forw,:soft,:back,:char); _model[k] = map(KnetArray, m[k]); end
    end
end

wsoft(model)=model[:soft][1]
bsoft(model)=model[:soft][2]
wchar(model)=model[:char][1]
bchar(model)=model[:char][2]
echar(model)=model[:cembed]
wforw(model)=model[:forw][1]
bforw(model)=model[:forw][2]
wback(model)=model[:back][1]
bback(model)=model[:back][2]

function minibatch(corpus, batchsize)
    data = Any[]
    sorted = sort(corpus, by=length)
    for i in 1:batchsize:length(corpus)
        j = min(length(corpus), i+batchsize-1)
        push!(data, sorted[i:j])
    end
    shuffle(data)
end

function minibatch1(corpus, batchsize)
    data = Any[]
    dict = Dict{Int,Vector{Int}}()
    for i in 1:length(corpus)
        s = corpus[i]
        l = length(s)
        a = get!(dict, l, Int[])
        push!(a, i)
        if length(a) == batchsize
            push!(data, corpus[a])
            empty!(a)
        end
    end
    for (k,a) in dict
        if !isempty(a)
            push!(data, corpus[a])
        end
    end
    return data
end

function main(;batch=3000,nsent=0,result=zeros(2))
    loadmodel()
    loadvocab()
    loadcorpus(_vocab)
    c = (nsent == 0 ? _corpus : _corpus[1:nsent])
    global _data = minibatch(c, batch)
    # global _data = [ _corpus[i:min(i+999,length(_corpus))] for i=1:1000:length(_corpus) ]
    # global _data = [ _corpus[1:1] ]
    for d in _data
        lmloss(_model, d, _vocab; result=result)
    end
    println(exp(-result[1]/result[2]))
    return result
end

function getloss(d; result=zeros(2))
    lmloss(_model, d, _vocab; result=result)
    println(exp(-result[1]/result[2]))
    return result
end

# Universal POS tags (17)
const UPOSTAG = Dict{String,PosTag}(
"ADJ"   => 1, # adjective
"ADP"   => 2, # adposition
"ADV"   => 3, # adverb
"AUX"   => 4, # auxiliary
"CCONJ" => 5, # coordinating conjunction
"DET"   => 6, # determiner
"INTJ"  => 7, # interjection
"NOUN"  => 8, # noun
"NUM"   => 9, # numeral
"PART"  => 10, # particle
"PRON"  => 11, # pronoun
"PROPN" => 12, # proper noun
"PUNCT" => 13, # punctuation
"SCONJ" => 14, # subordinating conjunction
"SYM"   => 15, # symbol
"VERB"  => 16, # verb
"X"     => 17, # other
]

# Universal Dependency Relations (37)
const UDEPREL = Dict{String,DepRel}(
"root"       => 1,  # root
"acl"        => 2,  # clausal modifier of noun (adjectival clause)
"advcl"      => 3,  # adverbial clause modifier
"advmod"     => 4,  # adverbial modifier
"amod"       => 5,  # adjectival modifier
"appos"      => 6,  # appositional modifier
"aux"        => 7,  # auxiliary
"case"       => 8,  # case marking
"cc"         => 9,  # coordinating conjunction
"ccomp"      => 10, # clausal complement
"clf"        => 11, # classifier
"compound"   => 12, # compound
"conj"       => 13, # conjunct
"cop"        => 14, # copula
"csubj"      => 15, # clausal subject
"dep"        => 16, # unspecified dependency
"det"        => 17, # determiner
"discourse"  => 18, # discourse element
"dislocated" => 19, # dislocated elements
"expl"       => 20, # expletive
"fixed"      => 21, # fixed multiword expression
"flat"       => 22, # flat multiword expression
"goeswith"   => 23, # goes with
"iobj"       => 24, # indirect object
"list"       => 25, # list
"mark"       => 26, # marker
"nmod"       => 27, # nominal modifier
"nsubj"      => 28, # nominal subject
"nummod"     => 29, # numeric modifier
"obj"        => 30, # object
"obl"        => 31, # oblique nominal
"orphan"     => 32, # orphan
"parataxis"  => 33, # parataxis
"punct"      => 34, # punctuation
"reparandum" => 35, # overridden disfluency
"vocative"   => 36, # vocative
"xcomp"      => 37, # open clausal complement
]

# """
# The input is minibatched and processed character based.
# The blstm language model used for preinit inputs characters outputs words.
# The characters of a word is used to obtain its embedding using an RNN or CNN.
# The input to the parser is minibatched sentences.
# The sentence lengths in a minibatch do not have to exactly match, we will support padding.
# """
# foo1

# """
# Let B be the minibatch size, and T the sentence length in the minibatch.
# input[t][b] is a string specifying the t'th word of the b'th sentence.
# Do we include SOS/EOS tokens in input?  No, the parser does not need them.
# How about the ROOT token?  No, our parsers are written s.t. the root token is implicit.
# """
# foo2

# """
# We need head and deprel for loss calculation. It makes more sense to
# process an array of sentences, doing the grouping inside.
# """
# foo3


text:
struct NonlinearSystem <: AbstractSystem
    eqs::Vector{Equation}
    vs::Vector{Variable}
    ps::Vector{Variable}
end

function NonlinearSystem(eqs)
    vs, ps = extract_elements(eqs, [_subtype(:Unknown), _subtype(:Parameter)])
    NonlinearSystem(eqs, vs, ps)
end

function generate_nlsys_function(sys::NonlinearSystem)
    var_exprs = [:($(sys.vs[i].name) = u[$i]) for i in 1:length(sys.vs)]
    param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in 1:length(sys.ps)]
    sys_eqs, calc_eqs = partition(eq -> isequal(eq.lhs, Constant(0)), sys.eqs)
    calc_exprs = [:($(eq.lhs.name) = $(eq.rhs)) for eq in calc_eqs if isa(eq.lhs, Variable)]
    sys_exprs = [:($(Symbol("resid[$i]")) = $(sys_eqs[i].rhs)) for i in eachindex(sys_eqs)]

    exprs = vcat(var_exprs,param_exprs,calc_exprs,sys_exprs)
    block = expr_arr_to_block(exprs)
    :((du,u,p)->$(block))
end

function calculate_jacobian(sys::NonlinearSystem,simplify=true)
    sys_eqs, calc_eqs = partition(eq -> isequal(eq.lhs, Constant(0)), sys.eqs)
    rhs = [eq.rhs for eq in sys_eqs]

    for calc_eq ∈ calc_eqs
        find_replace!.(rhs, calc_eq.lhs, calc_eq.rhs)
    end

    sys_exprs = calculate_jacobian(rhs,sys.vs)
    sys_exprs = Expression[expand_derivatives(expr) for expr in sys_exprs]
    sys_exprs
end

function generate_nlsys_jacobian(sys::NonlinearSystem,simplify=true)
    var_exprs = [:($(sys.vs[i].name) = u[$i]) for i in 1:length(sys.vs)]
    param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in 1:length(sys.ps)]
    jac = calculate_jacobian(sys,simplify)
    jac_exprs = [:(J[$i,$j] = $(convert(Expr, jac[i,j]))) for i in 1:size(jac,1), j in 1:size(jac,2)]
    exprs = vcat(var_exprs,param_exprs,vec(jac_exprs))
    block = expr_arr_to_block(exprs)
    :((J,u,p,t)->$(block))
end

export NonlinearSystem
export generate_nlsys_function

ext: jl


text:
# test spectral parameters for the filters
@testset "Spectral filter parameters" begin
# test ξmax
Q=12 # number of wavelets per octave
ξ=Scat.ξmax(Q)
@test ξ==1/(1+2^(0.25))
Q=1 # number of wavelets per octave
ξ=Scat.ξmax(Q)
@test ξ==0.35

# test σmin
J=1
σ0=0.1
σ = Scat.σmin(σ0, J)
@test σ==σ0/2.0

# test σmax
ξ = 1.0
rψ = 2.718281828459045^(-1/2)
Q = 1
σ = Scat.σmax(ξ, Q, rψ)
@test σ==1/3

# test dyadic subsampling
α=1.0
ξ=1.0
σ=1.0
mds = Scat.max_dyadic_subsampling(ξ, σ, α)
@test mds==0
α=0.0
ξ=0.25
σ=1.0
mds = Scat.max_dyadic_subsampling(ξ, σ, α)
@test mds==1

σ_min = 0.5
σ_max = 2.0
Q = 1
nds = Scat.n_dyadic_steps(Q, σ_min, σ_max)
@test nds==2

# test filter bank parameters
# single Q
Q=1
σ_min=0.05
rψ=sqrt(0.5)
α=5.0
σs, ξs, js = get_filter_params(Q, σ_min, rψ=rψ, α=α)
ξ_max = Scat.ξmax(Q)
σ_max = Scat.σmax(ξ_max, Q, rψ)
nds = Scat.n_dyadic_steps(Q, σ_min, σ_max)
collect(0.0:nds+Q-2)
@test length(σs)==nds+Q-1
@test length(ξs)==nds+Q-1
@test length(js)==nds+Q-1
# array of Q
Q=[12, 1]
J=12
σ0=0.1
rψ=sqrt(0.5)
α=5.0
σs, ξs, js = get_filter_params(Q, J, σ0, rψ=rψ, α=α)
@test length(σs)==length(Q)
@test length(ξs)==length(Q)
@test length(js)==length(Q)
end

# test FilterBank1d struct
@testset "Morlet wavelet filter bank" begin
# test inner constructor
N=12
Q=12
J=8
σ0=0.1
rψ=sqrt(0.5)
α=5.0

F = FilterBank1d(N, Q, J, σ0, rψ=rψ, α=α)
# test frequency vector
ω=FFTW.fftfreq(2^N)
@test ω==F.ω
# test number of wavelets in the filter bank
σs, ξs, js = get_filter_params(Q, σ0/2^J, rψ=rψ, α=α)
@test length(F.Λ)==length(σs)
end

# test get_FilterBanks
@testset "FilterBank1d convenience methods" begin
N=12
Q=[12,1]
J=[8,4]
σ0=[0.5, 0.1]
FilterBanks = get_FilterBanks(N,Q,J,σ0)
@test all([typeof(FilterBanks[i])==FilterBank1d for i=1:length(Q)])
@test length(FilterBanks)==length(Q)
@test (FilterBanks[1].Q==Q[1]) & (FilterBanks[2].Q==Q[2])
@test (FilterBanks[1].J==J[1]) & (FilterBanks[2].J==J[2])

# common J and σ0
N=12
Q=[12,1]
J=8
σ0=0.1
FilterBanks = get_FilterBanks(N,Q,J,σ0)
@test all([typeof(FilterBanks[i])==FilterBank1d for i=1:length(Q)])
@test length(FilterBanks)==length(Q)
@test (FilterBanks[1].Q==Q[1]) & (FilterBanks[2].Q==Q[2])
@test (FilterBanks[1].J==J) & (FilterBanks[2].J==J)

# common J
N=12
Q=[12,1]
J=8
σ0=[0.5, 0.1]
FilterBanks = get_FilterBanks(N,Q,J,σ0)
@test all([typeof(FilterBanks[i])==FilterBank1d for i=1:length(Q)])
@test length(FilterBanks)==length(Q)
@test (FilterBanks[1].Q==Q[1]) & (FilterBanks[2].Q==Q[2])
@test (FilterBanks[1].J==J) & (FilterBanks[2].J==J)

# common σ0
N=12
Q=[12,1]
J=[8, 4]
σ0=0.1
FilterBanks = get_FilterBanks(N,Q,J,σ0)
@test all([typeof(FilterBanks[i])==FilterBank1d for i=1:length(Q)])
@test length(FilterBanks)==length(Q)
@test (FilterBanks[1].Q==Q[1]) & (FilterBanks[2].Q==Q[2])
@test (FilterBanks[1].J==J[1]) & (FilterBanks[2].J==J[2])
end

# test FilterBank1dBlock struct
@testset "Filter bank vs block" begin
# test inner constructor
N=12
Q=12
J=8
σ0=0.1
F = FilterBank1d(N, Q, J, σ0)
Fblock = FilterBank1dBlock(N, Q, J, σ0)
# TODO: Find a better way to test the equality of the FilterBank1d and FilterBank1dBlock fields
testBoolΛ=true
testBoolσ=true
testBoolξ=true
testBoolj=true
for i=1:length(F.Λ)
    if F.Λ[i].ψ != Fblock.Λ[:,i]
        testBoolΛ=false
    end
    if F.Λ[i].σ != Fblock.σs[i]
        testBoolσ=false
    end
    if F.Λ[i].ξ != Fblock.ξs[i]
        testBoolξ=false
    end
    if F.Λ[i].j != Fblock.js[i]
        testBoolj=false
    end
end
# test whether morlet wavelets are the same
@test testBoolΛ=true
@test testBoolσ=true
@test testBoolξ=true
@test testBoolj=true
# test whether gaussian filter is the same
@test F.ϕ.ϕ==Fblock.ϕ
@test F.ϕ.σ==Fblock.σg
end

# test FilterBank1dBlock struct
@testset "FilterBank1dBlock convenience methods" begin
# common J and σ0
Fblocks = get_FilterBanks(12, [12, 1], 8, 0.1, typ=FilterBank1dBlock)
@test all([typeof(Fblocks[i])==FilterBank1dBlock for i=1:length(Q)])
@test length(Fblocks)==2
@test (Fblocks[1].Q==12) & (Fblocks[2].Q==1)
@test (Fblocks[1].J==8) & (Fblocks[2].J==8)

# common J
Fblocks = get_FilterBanks(12, [12,1], 8, [0.1, 0.5], typ=FilterBank1dBlock)
@test all([typeof(Fblocks[i])==FilterBank1dBlock for i=1:length(Q)])
@test length(Fblocks)==2
@test (Fblocks[1].Q==12) & (Fblocks[2].Q==1)
@test (Fblocks[1].J==8) & (Fblocks[2].J==8)

# common σ0
Fblocks = get_FilterBanks(12, [12,1], [8,4], 0.1, typ=FilterBank1dBlock)
@test all([typeof(Fblocks[i])==FilterBank1dBlock for i=1:length(Q)])
@test length(Fblocks)==2
@test (Fblocks[1].Q==12) & (Fblocks[2].Q==1)
@test (Fblocks[1].J==8) & (Fblocks[2].J==4)
end


text:
module GeostrophicThermalShallowWaterExplicit

using Gridap
using GridapGeosciences

include("ThermoGeostrophicInitialConditions.jl")

# Thermogeostrophic configuration using the explicit thermal shallow water
# solver.

order  = 1
degree = 4

n      = 24
nstep  = 20*24*30 # 20 days
dt     = 120.0

model = CubedSphereDiscreteModel(n; radius=rₑ)

hf, uf = thermal_shallow_water_mat_adv_explicit_time_stepper(model, order, degree,
                                    h₀, u₀, s₀, f,
                                    dt, 0.5*dt, 0.0*dt, nstep;
                                    write_solution=true,
                                    write_solution_freq=360,
                                    write_diagnostics=true,
                                    write_diagnostics_freq=1,
                                    dump_diagnostics_on_screen=true)

end

max_stars_repo_name: santiagobadia/GridapGeosciences

max_issues_repo_name: santiagobadia/GridapGeosciences

max_forks_repo_name: santiagobadia/GridapGeosciences


text:
module OpticalBlochEquations

using StaticArrays
using StructArrays
using Parameters
using HalfIntegers
using Unitful
using LinearAlgebra
using WignerSymbols
using LoopVectorization

macro with_unit(arg1, arg2)
    arg2 = @eval @u_str $arg2
    return convert(Float64, upreferred(eval(arg1) .* arg2).val)
end
export @with_unit

macro params(fields_tuple)
    fields = fields_tuple.args
    esc(
        quote
            NamedTuple{Tuple($fields)}(($fields_tuple))
        end)
end
export @params

const sph2cart = @SMatrix [
    -1/√2 +1/√2 0;
    +im/√2 +im/√2 0;
    0 0 1;
]
export sph2cart

const ϵ₊ = @SVector [-1/√2, -im/√2, 0]
const ϵ₋ = @SVector [ 1/√2, -im/√2, 0]
const ϵ₀ = @SVector [    0,      0, 1]
const ϵ = [ϵ₋, ϵ₀, ϵ₊]
export ϵ₊, ϵ₋, ϵ₀
export ϵ

const qs = @SVector [-1, 0, 1]
export qs

x̂ = @SVector [1,0,0]
ŷ = @SVector [0,1,0]
ẑ = @SVector [0,0,1]
#export x̂, ŷ, ẑ

# function make_hamiltonian(states::Vector(State), lasers::Vector(Laser))
#     H = Hamiltonian(; states)
#     return H
# end

@with_kw struct Laser1
    k::SVector{3, Float64}      # k-vector
    e::SVector{3, ComplexF64}   # polarization
    ω::Float64                  # frequency
    s::Float64                  # saturation parameter
    fre::SVector{3, Float64}
    fim::SVector{3, Float64}
    Hq::Vector{Array{Float64, 2}}
end

@with_kw struct Field
    k::SVector{3, Float64}      # k-vector
    e::SVector{3, ComplexF64}   # polarization
    ω::Float64                  # frequency
    s::Float64                  # saturation parameter
    fre::SVector{3, Float64}
    fim::SVector{3, Float64}
    Hqm::Array{Float64, 2}
    Hq0::Array{Float64, 2}
    Hqp::Array{Float64, 2}
end

@with_kw struct State
    F::HalfInt                          # angular quantum number
    m::HalfInt                          # projection of angular quantum number
    ω::Float64                          # frequency
    μ::Float64                          # magnetic moment
    Γ::Union{Nothing, Float64}=nothing  # linewidth (defaults to `nothing` for ground states)
end
export State

@with_kw struct Hamiltonian12
    states::Vector{State}
    H::StructArray{<:Complex} = StructArray(zeros(ComplexF64, length(states), length(states)))
    Hsubs::Vector{StructArray{<:Complex}} = []
end
export Hamiltonian12

function define_laser(k, e, ω, s)
    fre = zeros(Float64, 3)
    fim = zeros(Float64, 3)
    for q in eachindex(qs)
        dotted = e ⋅ ϵ[q]
        fre[q] = real(dotted)
        fim[q] = imag(dotted)
    end
    Hq = [zeros(Float64, (12, 12)) for _ in qs]
    return Laser1(k, e, ω, s, fre, fim, Hq)
end
export define_laser

function define_field(k, e, ω, s)
    fre = zeros(Float64, 3)
    fim = zeros(Float64, 3)
    for q in eachindex(qs)
        dotted = e ⋅ ϵ[q]
        fre[q] = real(dotted)
        fim[q] = imag(dotted)
    end
    Hqm = zeros(Float64, (3, 3))
    Hq0 = zeros(Float64, (3, 3))
    Hqp = zeros(Float64, (3, 3))
    return Field(k, e, ω, s, fre, fim, Hqm, Hq0, Hqp)
end
export define_field

@with_kw struct Manifold
    F::HalfInt                          # angular quantum number
    states::Vector{State}
end
manifold(; F, ω, μ, Γ=nothing) = Manifold(F, [State(F=F, m=m, ω=ω, μ=μ, Γ=Γ) for m in -F:F])
export manifold

# Structure for quantum jumps from state `s` to state `s′` with rate `r`
struct Jump
    s ::Int64
    s′::Int64
    r ::Float64
end

roundmult(val, prec) = (inv_prec = 1 / prec; round(val * inv_prec) / inv_prec)
function round_freq(ω, Γ)
    ω_min = Γ * 1e-2
    return roundmult(ω, ω_min)
end

function schrödinger(states, lasers, d, f, p)

    n_states = length(states)
    n_lasers = length(lasers)

    lasers = StructArray(lasers)

    r = @SVector [0.,0.,0.]
    v = @SVector [0.,0.,0.]

    type_complex = ComplexF64
    type_real    = Float64

    H = StructArray( zeros(type_complex, n_states, n_states) )

    # Define the optical Hamiltonian; it has dimensions (k, q)
    for i ∈ 1:n_states, j ∈ i:n_states
        for k in eachindex(lasers)
            lasers[k].Hqm[i,j] += d[1,i,j] * (lasers[k].e ⋅ ϵ[1])
            lasers[k].Hq0[i,j] += d[2,i,j] * (lasers[k].e ⋅ ϵ[2])
            lasers[k].Hqp[i,j] += d[3,i,j] * (lasers[k].e ⋅ ϵ[3])
            lasers[k].Hqm[j,i] += d[1,i,j] * (lasers[k].e ⋅ ϵ[1])
            lasers[k].Hq0[j,i] += d[2,i,j] * (lasers[k].e ⋅ ϵ[2])
            lasers[k].Hqp[j,i] += d[3,i,j] * (lasers[k].e ⋅ ϵ[3])
        end
    end

    ψ = zeros(type_complex, n_states)
    dψ = deepcopy(ψ)
    ψ[1] = 1
    ψ_soa = StructArray(ψ)
    dψ_soa = StructArray(dψ)

    p = @params (H, ψ_soa, dψ_soa, states, lasers, r, v, f, p)

    return (dψ, ψ, p)
end
export schrödinger

function obe(states, lasers, d, ρ)

    n_states = length(states)
    n_lasers = length(lasers)

    lasers = StructArray(lasers)

    r = @SVector [0., 0., 0.]
    v = @SVector [0., 0., 0.]

    type_complex = ComplexF64
    type_real = Float64

    H       = StructArray( zeros(type_complex, n_states, n_states) )
    H_adj   = StructArray( zeros(type_complex, n_states, n_states) )
    # HJ      = zeros(type_real, n_states, n_states)
    HJ      = zeros(Float64, n_states)
    Hₒ      = [zeros(type_real, n_states, n_states) for l in lasers, q in qs]

    # Define the optical Hamiltonian; it has dimensions (k, q)
    for k in eachindex(lasers)
        for s in eachindex(states), s′ in s:n_states
            q = Int64(states[s′].m - states[s].m)
            if abs(q) <= 1
                if !iszero(lasers[k].e ⋅ ϵ[q+2])
                    # Hₒ[k, q+2][s, s′] = d[s, s′, q+2]
                    # Hₒ[k, q+2][s′, s] = d[s, s′, q+2]
                    lasers[k].Hq[q+2][s, s′] = d[s, s′, q+2]
                    lasers[k].Hq[q+2][s′, s] = d[s, s′, q+2]
                end
            end
        end
    end

    # Define the magnetic Hamiltonian; it has dimensions (q)
    # μ = (s, s′, q) -> - s.g * (-1)^(s.F - s′.m) * sqrt(s.F * (s.F + 1) * (2s.F + 1))
    #     * wigner3j(s.F, 1, s′.F, -s.m, q, s′.m)
    # for s in eachindex(states), s′ in eachindex(states), q in qs
    #     if states[s].F == states[s].F # States only mix if they belong to the same F state
    #         Hₘ[q+2][s, s′] = im * (-1)^q * μ(states[s], states[s′], -q)
    #     end
    # end

    # Construct an array containing all jump operators, as defined by `d`
    Js = Array{Jump}(undef, 0)
    for s′ in eachindex(states), s in s′:n_states, q in qs
        dme = d[s, s′, q+2]
        if dme != 0 & (states[s′].ω < states[s].ω) # only energy-allowed jumps are generated
            J = Jump(s, s′, dme)
            push!(Js, J)
        end
    end

    for J in Js
        # Adds the term - (iħ / 2) ∑ᵢ Jᵢ† Jᵢ
        # We assume jumps take the form Jᵢ = sqrt(Γ)|g⟩⟨e| such that Jᵢ† Jᵢ = Γ|e⟩⟨e|
        HJ[J.s] -= 0.5 * J.r^2
    end

    ω = [s.ω for s in states]
    eiωt  = StructArray(zeros(type_complex, n_states))

    dρ = deepcopy(ρ)
    ρ_soa = StructArray(ρ)
    dρ_soa = StructArray(dρ)

    # B(r) = [0,0,0]

    Γs = [s.Γ for s in states]
    Γ = maximum(filter(x -> x != nothing, Γs))

    conj_mat = ones(Float64, n_states, n_states)
    for i in 1:n_states, j in 1:n_states
        if j < i
            conj_mat[i,j] = -1
        end
    end

    # Allocate some temporary arrays
    A12 = zeros(type_real, n_states, n_states)
    B12 = zeros(type_real, n_states, n_states)
    T1 = zeros(type_real, n_states, n_states)
    T2 = zeros(type_real, n_states, n_states)
    tmp1 = StructArray(zeros(ComplexF64, n_states, n_states))
    tmp2 = StructArray(zeros(ComplexF64, n_states, n_states))

    p = @params (H, HJ, ρ_soa, dρ_soa, Js, ω, eiωt, states, lasers, r, v, Γ, conj_mat, Hₒ, A12, B12, T1, T2, tmp1, tmp2)

    return (dρ, ρ, p)
end
export obe

function update_H_schrödinger!(τ, v, lasers, H)

    @turbo for l in eachindex(lasers)
        s = lasers.s[l]
        Ho_laser_m = lasers.Hqm[l]
        Ho_laser_0 = lasers.Hq0[l]
        Ho_laser_p = lasers.Hqp[l]
        for i in eachindex(H)
            H.re[i] += s * (Ho_laser_m[i] + Ho_laser_0[i] + Ho_laser_p[i])
        end
    end

    return nothing
end

# function force(lasers, Γ)
#
#     x = - h * Γ / (2 * √2 )
#     for l in eachindex(lasers)

function update_H!(τ, v, Γ, lasers, H, conj_mat)

    @turbo H.re .= 0
    @turbo H.im .= 0
    r = v .* (τ / Γ)

    for l in eachindex(lasers)
        s = lasers.s[l]
        k = lasers.k[l]
        e = lasers.e[l]
        ω = lasers.ω[l]
        fre = lasers.fre[l]
        fim = lasers.fim[l]
        Ho_laser = lasers.Hq[l]
        c, v = sincos(k ⋅ r - ω * τ)
        x = sqrt(s) / (2 * √2)
        for q in eachindex(qs)
            freq = x * (fre[q] * v - fim[q] * c)
            fimq = x * (fre[q] * c + fim[q] * v)
            Ho_laser_q = Ho_laser[q]
            if (freq > 1e-10) || (fimq > 1e-10) || (freq < -1e-10) || (freq < -1e-10)
                @turbo for i in eachindex(H)
                    a = Ho_laser_q[i]
                    H.re[i] += freq * a
                    H.im[i] += conj_mat[i] * fimq * a
                end
            end
        end
    end

    # @inbounds @simd for q in qs
    #     β = (μB / (Γ * ħ)) * (B(r) ⋅ ϵ[q+2])
    #     axpy!(β, Hₘ[q+2], H)
    # end

    return nothing
end

function soa_to_base!(ρ::Array{<:Complex}, ρ_soa::StructArray{<:Complex})
    @turbo for i in eachindex(ρ, ρ_soa)
        ρ[i] = ρ_soa.re[i] + im * ρ_soa.im[i]
    end
    return nothing
end

function base_to_soa!(ρ::Array{<:Complex}, ρ_soa::StructArray{<:Complex})
    @turbo for i in eachindex(ρ, ρ_soa)
        ρ_soa.re[i] = real(ρ[i])
        ρ_soa.im[i] = imag(ρ[i])
    end
    return nothing
end

# function base_to_base!(ρ::Array{<:Complex}, ρ_soa::StructArray{<:Complex})
#     @turbo for i in eachindex(ρ, ρ_soa)
#         ρ[i] = ρ_soa.re[i] + im * ρ_soa.im[i]
#     end
#     return nothing
# end
#
# function soa_to_base_mult!(ρ::Array{<:Complex}, ρ_soa::StructArray{<:Complex})
#     @turbo for i in eachindex(ρ, ρ_soa)
#         ρ[i] *= ρ_soa.re[i] + im * ρ_soa.im[i]
#     end
#     return nothing
# end

function update_eiωt!(eiωt::StructArray{<:Complex}, ω::Array{<:Real}, τ::Real)
    @turbo for i ∈ 1:size(ω, 1)
        eiωt.im[i], eiωt.re[i] = sincos( ω[i] * τ )
    end
    return nothing
end

function Heisenberg!(ρ::StructArray{<:Complex}, eiωt::StructArray{<:Complex}, im_factor=1)
    @turbo for j ∈ 1:size(ρ, 2)
        jre = eiωt.re[j]
        jim = eiωt.im[j]
        for i ∈ 1:size(ρ, 1)
            ire = eiωt.re[i]
            iim = eiωt.im[i]
            cisim = im_factor * (iim * jre - ire * jim)
            cisre = ire * jre + iim * jim
            ρre_i = ρ.re[i,j]
            ρim_i = ρ.im[i,j]
            ρ.re[i,j] = ρre_i * cisre - ρim_i * cisim
            ρ.im[i,j] = ρre_i * cisim + ρim_i * cisre
        end
    end
    return nothing
end
export Heisenberg!

#
# function mul_times!(A::StructArray{<:Complex}, B::StructArray{<:Complex}, im_factor=1)
#     @turbo for i ∈ eachindex(A)
#         A.re[i] *= B.re[i]
#         A.im[i] *= im_factor * B.im[i]
#     end
# end

function im_commutator!(C, A, B, tmp1, tmp2, A_diag)
    @turbo for i ∈ 1:size(A,1), j ∈ 1:size(B,2)
        Cre = 0.0
        Cim = 0.0
        for k ∈ 1:size(A,2)
            Aik_re = A.re[i,k]
            Aik_im = A.im[i,k]
            Bkj_re = B.re[k,j]
            Bkj_im = B.im[k,j]
            # Multiply by -1
            Cre -= Aik_re * Bkj_re - Aik_im * Bkj_im
            Cim -= Aik_re * Bkj_im + Aik_im * Bkj_re
        end
        C.re[i,j] = Cre
        C.im[i,j] = Cim
    end

    adjoint!(tmp1, C)
    C_add_A!(C, tmp1, -1)

    mul_by_im!(C)

    mul_diagonal!(tmp1, B, A_diag)
    adjoint!(tmp2, tmp1)

    C_add_AplusB!(C, tmp1, tmp2, 1, 1)

end
export im_commutator!

# function im_commutator_old!(C::StructArray{<:Complex}, A::StructArray{<:Complex}, B::StructArray{<:Complex}, A12::Array{<:Real}, B12::Array{<:Real}, T1::Array{<:Real}, T2::Array{<:Real}, A_diag::Array{<:Real}, tmp1::StructArray{<:Complex}, tmp2::StructArray{<:Complex})
#     """
#     Computes `-im * [A, B] = -im * (A * B - B * A)` in-place by overwriting `C`.
#
#     3m complex multiplication (includes a factor -1):
#     -C_re = T2 - T1
#     -C_im = T1 + T2 - (A1 + A2)(B1 + B2)
#
#     (Should also try via a single matrix multiplication defined on a StructArray type.)
#     """
#
#     update_T₁T₂!(T1, T2, A, B)
#     C_copy_AplusB!(A12, A.re, A.im)
#     C_copy_AplusB!(B12, B.re, B.im)
#
#     C_copy_T₁T₂!(C, T1, T2)
#     #jmul!(C.im, A12, B12, -1, 1)
#
#     adjoint!(tmp1, C)
#     C_add_A!(C, tmp1, -1)
#
#     mul_by_im!(C)
#
#     mul_diagonal!(tmp1, B, A_diag)
#     adjoint!(tmp2, tmp1)
#
#     C_add_AplusB!(C, tmp1, tmp2, 1, 1)
#
#     return nothing
# end

# function jgemvavx!(y, A, x)
#     @turbo for i ∈ eachindex(y)
#         yi_re = 0.0
#         yi_im = 0.0
#         for j ∈ eachindex(x)
#             xj_re = x.re[j]
#             xj_im = x.im[j]
#             Aij_re = A.re[i,j]
#             Aij_im = A.im[i,j]
#             yi_re += Aij_re * xj_re - Aij_im * xj_im
#             yi_im -= Aij_re * xj_im + Aij_im * xj_re
#         end
#         y.re[i] = yi_im
#         y.im[i] = yi_re
#     end
#     return y
# end

function ψ!(dψ, ψ, p, τ)

    # @avx p.H .= 0
    # @avx p.ψ_soa .= ψ
    base_to_soa!(ψ, p.ψ_soa)
    # Update the Hamiltonian according to the new time τ
    p.f(p.H, τ, p.p)
    # update_H_schrödinger!(τ, p.v, p.lasers, p.H)
    jgemvavx!(p.dψ_soa, p.H, p.ψ_soa)
    soa_to_base!(dψ, p.dψ_soa)

    return nothing
end
export ψ!

function ρ!(dρ, ρ, p, τ)

    base_to_soa!(ρ, p.ρ_soa)
    #p.ρ_soa .= ρ

    # Update the Hamiltonian according to the new time τ
    update_H!(τ, p.v, p.Γ, p.lasers, p.H, p.conj_mat)

    # Apply a transformation to go to the Heisenberg picture
    update_eiωt!(p.eiωt, p.ω, τ)
    Heisenberg!(p.ρ_soa, p.eiωt)

    # Compute coherent evolution terms
    # im_commutator!(p.dρ_soa, p.H, p.ρ_soa, p.A12, p.B12, p.T1, p.T2, p.HJ, p.tmp1, p.tmp2)
    im_commutator!(p.dρ_soa, p.H, p.ρ_soa, p.tmp1, p.tmp2, p.HJ)

    # Add the terms ∑ᵢ Jᵢ ρ Jᵢ†
    # We assume jumps take the form Jᵢ = sqrt(Γ)|g⟩⟨e| such that JᵢρJᵢ† = Γ^2|g⟩⟨g|ρₑₑ
    @inbounds for i in eachindex(p.Js)
        J = p.Js[i]
        p.dρ_soa.re[J.s′, J.s′] += J.r^2 * p.ρ_soa.re[J.s, J.s]
    end

    # The left-hand side also needs to be transformed into the Heisenberg picture
    # To do this, we require the transpose of the `ω` matrix
    # Heisenberg!(p.dρ_soa, p.ω_trans, τ)
    Heisenberg!(p.dρ_soa, p.eiωt, -1)
    soa_to_base!(dρ, p.dρ_soa)

    return nothing
end
export ρ!

function C_copy_AplusB!(C::Array{<:Real}, A::Array{<:Real}, B::Array{<:Real}, α=1, β=1)
    @turbo for i in eachindex(A, B, C)
        C[i] = α * A[i] + β * B[i]
    end
end

function C_add_AplusB!(C::Array{<:Real}, A::Array{<:Real}, B::Array{<:Real}, α=1, β=1)
    @turbo for i in eachindex(A, B, C)
        C[i] += α * A[i] + β * B[i]
    end
end

function C_add_AplusB!(C::StructArray{<:Complex}, A::StructArray{<:Complex}, B::StructArray{<:Complex},
    α=1, β=1)
    @turbo for i in eachindex(A, B, C)
        C.re[i] += α * A.re[i] + β * B.re[i]
        C.im[i] += α * A.im[i] + β * B.im[i]
    end
end

function C_add_A!(C::StructArray{<:Complex}, A::StructArray{<:Complex}, factor=1)
    @turbo for i in eachindex(C, A)
        C.re[i] += factor * A.re[i]
        C.im[i] += factor * A.im[i]
    end
end

function C_copy_T₁T₂!(C::StructArray{<:Complex}, T1::Array{<:Real}, T2::Array{<:Real})
    @turbo for i ∈ eachindex(C)
        C.re[i] = T2[i] - T1[i]
        C.im[i] = T2[i] + T1[i]
    end
end

function mul_by_im!(C::StructArray{<:Complex})
    @turbo for i ∈ eachindex(C)
        a = C.re[i]
        C.re[i] = -C.im[i]
        C.im[i] = a
    end
end

function mul_diagonal!(C::StructArray{<:Complex}, A::StructArray{<:Complex}, B::Array{<:Real})
    """
    Computes A × B, where B is diagonal and real.
    """
    @turbo for j in axes(A,2)
        d = B[j]
        for i in axes(A,1)
            C.re[i,j] = d * A.re[i,j]
            C.im[i,j] = d * A.im[i,j]
        end
    end
end

function update_T₁T₂!(T1::Array{<:Real}, T2::Array{<:Real}, A::StructArray{<:Complex}, B::StructArray{<:Complex})
    @turbo for i ∈ 1:size(A,1), j ∈ 1:size(B,2)
        C1 = zero(eltype(A))
        C2 = zero(eltype(A))
        for k ∈ 1:size(A,2)
            C1 += A.re[i,k] * B.re[k,j]
            C2 += A.im[i,k] * B.im[k,j]
        end
        T1[i,j] = C1
        T2[i,j] = C2
    end
end

end


text:
# Uncomment below if the package is not already installed
#using Pkg; Pkg.add("DifferentialEquations")

using DifferentialEquations, PyPlot
f(y,p,t) = [y[2], -sin(y[1])]
y0 = [2.5, 0]

tspan = (0.0, 10.0)

prob = ODEProblem(f, y0, tspan)
sol = solve(prob, abstol=1e-8, reltol=1e-8);

display(hcat(sol.t, collect(sol.u)))

tt = 0:0.1:10
plot(tt, sol.(tt));
legend(("\$\\theta(t)\$", "\$\\theta'(t)\$"));

condition(y, t, integrator) = y[2]

affect!(integrator) = terminate!(integrator)

cb = ContinuousCallback(condition, affect!)
sol = solve(prob, callback=cb, abstol=1e-8, reltol=1e-8)
println("Period = $(2sol.t[end])")


text:
"""
    `scatter(x, y, z)` / `scatter(x, y)` / `scatter(positions)`

Plots a marker for each element in `(x, y, z)`, `(x, y)`, or `positions`.
"""
function draw(scene::Scene, plot::Scatter)
    # Extracts attributes from plot and makes them available in the current scope
    # This also applies all conversions + extracts the current value from them.
    # e.g. plot[:color] might be e.g. Node(:red), while @get_attribute(plot, (color,)) will be
    # RGBAf0(1, 0, 0, 1) in this case!
    @get_attribute(plot, (marker,
        markersize,strokecolor,strokewidth,glowcolor,glowwidth,rotations,
        colormap,colorrange,marker_offset,fxaa,transform_marker, # Applies the plots transformation to marker
        uv_offset_width, distancefield, model
    ))
    positions_node = plot[1] # a node with positions, that will update its value whenever changed by user
    #extract the current value - alternativel just use positions_node[] to extract! Will always be a vector of points...
    positions = to_value(positions_node)::AbstractVector{<: Point}

    projected_positions = map(positions) do pos
        # take camera from scene + model transformation matrix and apply it to pos
        Makie.CairoBackend.project_position(scene, pos, model)
    end
    projected_positions::Vector{Vec2f0} # now you get 2D positions correctly transformed by the camera to screen pixel coordinates
end

"""
    `image(x, y, image)` / `image(image)`

Plots an image on range `x, y` (defaults to dimensions).
"""
function draw(scene::Scene, plot::Image)
    @get_attribute plot (colormap, colorrange)
end

# could be implemented via image, but might be optimized specifically by the backend
"""
    `heatmap(x, y, values)` or `heatmap(values)`

Plots a heatmap as an image on `x, y` (defaults to interpretation as dimensions).
"""
function draw(scene::Scene, plot::Heatmap)
    @get_attribute(plot, (colormap,colorrange,linewidth,levels,fxaa,interpolate))
end

"""
    `volume(volume_data)`

Plots a volume. Available algorithms are:
* `:iso` => IsoValue
* `:absorption` => Absorption
* `:mip` => MaximumIntensityProjection
* `:absorptionrgba` => AbsorptionRGBA
* `:indexedabsorption` => IndexedAbsorptionRGBA
"""
function draw(scene::Scene, plot::Volume)
    @get_attribute(plot, (algorithm,absorption,isovalue,isorange,colormap,colorrange))

end

"""
    `surface(x, y, z)`

Plots a surface, where `(x, y, z)` are supposed to lie on a grid.
"""
function draw(scene::Scene, plot::Surface)
    @get_attribute(plot, (colormap,colorrange,shading))
end

"""
    `lines(x, y, z)` / `lines(x, y)` / or `lines(positions)`

Creates a connected line plot for each element in `(x, y, z)`, `(x, y)` or `positions`.
"""
function draw(scene::Scene, plot::Lines)
    @get_attribute(plot, (linewidth,colormap,colorrange,linestyle))
end

"""
    `linesegments(x, y, z)` / `linesegments(x, y)` / `linesegments(positions)`

Plots a line for each pair of points in `(x, y, z)`, `(x, y)`, or `positions`.

**Attributes**:
The same as for [`lines`](@ref)
"""
function draw(scene::Scene, plot::LineSegments)
    @get_attribute(plot, (linewidth,colormap,colorrange,linestyle))
end

# alternatively, mesh3d? Or having only mesh instead of poly + mesh and figure out 2d/3d via dispatch
"""
    `mesh(x, y, z)`, `mesh(mesh_object)`, `mesh(x, y, z, faces)`, or `mesh(xyz, faces)`

Plots a 3D mesh.
"""
function draw(scene::Scene, plot::Mesh)
    @get_attribute(plot, (interpolate,shading,colormap,colorrange))
end

"""
    `meshscatter(x, y, z)` / `meshscatter(x, y)` / `meshscatter(positions)`

Plots a mesh for each element in `(x, y, z)`, `(x, y)`, or `positions` (similar to `scatter`).
`markersize` is a scaling applied to the primitive passed as `marker`
"""
function draw(scene::Scene, plot::MeshScatter)
    @get_attribute(plot, (marker,markersize,rotations,colormap,colorrange))
end

"""
    `text(string)`

Plots a text.
"""
function draw(scene::Scene, plot::Text)
    @get_attribute(plot, (strokecolor,strokewidth,font,align,rotation,textsize,position))
end


text:
### Mean

@mkapprule Mean :nargs => 1
@doap Mean(x::ListT) = Statistics.mean(margs(x))


text:
__precompile__()

module NetalignUtils

using NetalignMeasures

include("types.jl")
include("utils.jl")
include("dynet.jl")
include("ionet.jl")
include("iodynet.jl")
include("iomeasures.jl")
include("ioaln.jl")
include("iomatrix.jl")
include("iogo.jl")
include("generators.jl")
include("randomize.jl")

end # module


text:
function ensemble_history(e::GMC_NS_Ensemble, bins=25)
    !e.sample_posterior && throw(ArgumentError("This ensemble has no posterior samples to show a history for!"))
    livec=vcat([model.log_Li for model in e.models],[model.log_Li for model in e.posterior_samples])
    show(histogram(livec, nbins=bins))
end

function e_backup(e::GMC_NS_Ensemble, tuner::Dict{Int64,<:GMC_Tuner})
    cp(string(e.path,'/',"ens"),string(e.path,'/',"ens.bak"), force=true)
    serialize(string(e.path,'/',"ens"), e)
    serialize(string(e.path,'/',"tuner"), tuner)
end

function clean_ensemble_dir!(e::GMC_NS_Ensemble)
    if e.sample_posterior
        ens_paths=vcat(
            [basename(m.path) for m in e.models],
            [basename(m.path) for m in e.posterior_samples],
            "ens", "tuner"
        )
    else
        ens_paths=vcat(
            [basename(m.path) for m in e.models],
            "ens", "tuner"
        )
    end

    for file in readdir(e.path)
        !(file in ens_paths) && rm(e.path*'/'*file)
    end
end

function complete_evidence(e::GMC_NS_Ensemble)
    log_Z=e.log_Zi[end]
    H=e.Hi[end]
    live_weight=lps(e.log_Xi[length(e.log_Li)], -log(length(e.models)))
    lis=sort([model.log_Li for model in e.models])
    liwis=sort(lps.(lis, live_weight))

    for (li,liwi) in zip(lis,liwis)
        last_Z=log_Z
        log_Z=logaddexp(log_Z,liwi)
        H=lps(
            (exp(lps(liwi,-log_Z)) * li),
            (exp(lps(last_Z,-log_Z)) * lps(H,last_Z)),
            -log_Z)
    end

    return log_Z, H
end

function measure_evidence(e::GMC_NS_Ensemble)
    log_Z, H = complete_evidence(e)
    return measurement(log_Z,sqrt(abs(H)/length(e.models)))
end

function reset_ensemble!(e::GMC_NS_Ensemble)
    Ni=Int64.(round.(collect(-1:-1:-length(e.log_Li)+1)./e.log_Xi[2:end]))
    N=Ni[1]

    new_models=Vector{typeof(e.models[1])}()

    for traj in 1:N
        if string(traj)*".1" in [basename(record.path) for record in e.models]
            push!(new_models,e.models[findfirst(isequal(string(traj)*".1"), [basename(record.path) for record in e.models])])
        else
            push!(new_models,e.posterior_samples[findfirst(isequal(string(traj)*".1"), [basename(record.path) for record in e.posterior_samples])])
        end
    end

    e.models=new_models
    e.posterior_samples=Vector{typeof(e.models[1])}()

    e.contour=minimum([record.log_Li for record in e.models])

    e.log_Li=[e.log_Li[1]]
    e.log_Xi=[e.log_Xi[1]]
    e.log_wi=[e.log_wi[1]]
    e.log_Liwi=[e.log_Liwi[1]]
    e.log_Zi=[e.log_Zi[1]]
    e.Hi=[e.Hi[1]]

    e.t_counter=length(e.models)+1

    clean_ensemble_dir!(e)
    isfile(e.path*"/tuner") && rm(e.path*"/tuner")
    serialize(e.path*"/ens", e)

    return e
end

function move_ensemble!(e::GMC_NS_Ensemble,path::String)
    !isdir(path) && mkpath(path)
    rectype=typeof(e.models[1])

    for file in readdir(e.path)
        mv(e.path*'/'*file,path*'/'*file)
    end

    for (n,model) in enumerate(e.models)
        e.models[n]=rectype(model.trajectory, model.i, model.pos, path*'/'*basename(model.path), model.log_Li)
    end
    if e.sample_posterior
        for (n,model) in enumerate(e.posterior_samples)
            e.posterior_samples[n]=rectype(model.trajectory, model.i, model.pos, path*'/'*basename(model.path), model.log_Li)
        end
    end

    rm(e.path)
    e.path=path
    serialize(e.path*"/ens",e)
    return e
end

function copy_ensemble(e::GMC_NS_Ensemble,path::String)
    new_e=deepcopy(e)
    !isdir(path) && mkdir(path)
    rectype=typeof(e.models[1])
    for file in readdir(e.path)
        cp(e.path*'/'*file,path*'/'*file)
    end

    for (n,model) in enumerate(e.models)
        new_e.models[n]=rectype(model.trajectory, model.i, model.pos, path*'/'*basename(model.path), model.log_Li)
    end
    if e.sample_posterior
        for (n,model) in enumerate(e.posterior_samples)
            new_e.posterior_samples[n]=rectype(model.trajectory, model.i, model.pos, path*'/'*basename(model.path), model.log_Li)
        end
    end

    new_e.path=path
    serialize(new_e.path*"/ens",new_e)
    return new_e
end

function rewind_ensemble!(e::GMC_NS_Ensemble,rewind_idx)
    !e.sample_posterior && throw(ArgumentError("An ensemble not retaining posterior samples cannot be rewound!"))
    rewind_idx >= length(e.log_Li) && throw(ArgumentError("rewind_idx must be less than the current iterate!"))

    Ni=Int64.(round.(collect(-1:-1:-length(e.log_Li)+1)./e.log_Xi[2:end]))
    insert!(Ni,1,Ni[1])

    for it in length(e.log_Li)-1:-1:rewind_idx
        last_post=pop!(e.posterior_samples)

        if Ni[it]==Ni[it-1]
            models_idx=findfirst(isequal(last_post.trajectory), [rec.trajectory for rec in e.models])
            insert!(e.models, models_idx, last_post)
            deleteat!(e.models, models_idx+1)
        else
            push!(e.models, last_post)
        end

        !(length(e.models)==Ni[it-1]) && throw(DomainError("it: $it Ni[it]:$(Ni[it]) Ni[it-1]:$(Ni[it-1])"))
    end

    e.contour=e.log_Li[rewind_idx]
    e.log_Li=e.log_Li[1:rewind_idx]
    e.log_Xi=e.log_Xi[1:rewind_idx]
    e.log_wi=e.log_wi[1:rewind_idx]
    e.log_Liwi=e.log_Liwi[1:rewind_idx]
    e.log_Zi=e.log_Zi[1:rewind_idx]
    e.Hi=e.Hi[1:rewind_idx]

    traj_nos=vcat([rec.trajectory for rec in e.models],[rec.trajectory for rec in e.posterior_samples])

    e.t_counter=maximum(traj_nos)

    return e
end

function show_models(e::GMC_NS_Ensemble,idxs)
    liperm=sortperm([model.log_Li for model in e.models],rev=true)
    for idx in idxs
        m=deserialize(e.models[liperm[idx]].path)
        show(m)
    end
end

function show_models_e(e::GMC_NS_Ensemble,idxs)
    liperm=sortperm([model.log_Li for model in e.models],rev=true)
    for idx in idxs
        m=deserialize(e.models[liperm[idx]].path)
        show(stdout,m,e)
    end
end

function get_model(e::GMC_NS_Ensemble,no)
    return deserialize(e.path*'/'*string(no))
end

function reestimate_ensemble!(e::GMC_NS_Ensemble, wi_mode="trapezoidal")
    Ni=Int64.(round.(collect(-1:-1:-length(e.log_Li)+1)./e.log_Xi[2:end]))
    push!(Ni, length(e.models))

    #fix any borked starting values
    e.log_Li[1]=-Inf #L0 = 0
	e.log_Xi[1]=0. #X0 = 1
	e.log_wi[1]=-Inf #w0 = 0
	e.log_Liwi[1]=-Inf #Liwi0 = 0
	e.log_Zi[1]=-Inf #Z0 = 0
	e.Hi[1]=0. #H0 = 0,

    for i in 1:length(e.log_Li)-1
        j=i+1

        e.log_Li[j]=e.posterior_samples[i].log_Li

        e.log_Xi[j]=-i/Ni[i]

        if wi_mode=="trapezoidal"
            e.log_wi[j]= logsubexp(e.log_Xi[i], -j/Ni[j]) - log(2) #log width of prior mass spanned by the last step-trapezoidal approx
        elseif wi_mode=="simple"
            e.log_wi[j]= logsubexp(e.log_Xi[i], e.log_Xi[j]) #log width of prior mass spanned by the last step-simple approx
        else
            throw(ArgumentError("Unsupported wi_mode!"))
        end
        e.log_Liwi[j]=lps(e.log_Li[j],e.log_wi[j]) #log likelihood + log width = increment of evidence spanned by iterate
        e.log_Zi[j]=logaddexp(e.log_Zi[i],e.log_Liwi[j])    #log evidence
        #information- dimensionless quantity
        Hj=lps(
            (exp(lps(e.log_Liwi[j],-e.log_Zi[j])) * e.log_Li[j]),
            (exp(lps(e.log_Zi[i],-e.log_Zi[j])) * lps(e.Hi[i],e.log_Zi[i])),
            -e.log_Zi[j])
        Hj === -Inf ? (e.Hi[j]=0.) : (e.Hi[j]=Hj)
    end
end

#get a posterior kernel density estimate from the ensemble
function posterior_kde(e::GMC_NS_Ensemble; bivar=Vector{Pair{<:Integer}}())
    !e.sample_posterior && throw(ArgumentError("This ensemble is not set to collect posterior samples!"))

    θ_mat=zeros(length(e.models[1].pos),length(e.models)+length(e.posterior_samples))
    weight_vec=zeros(length(e.models)+length(e.posterior_samples))

    logz, _=complete_evidence(e)

    for (n,mrec) in enumerate(e.posterior_samples)
        m=deserialize(mrec.path)
        θ_mat[:,n]=m.θ
        weight_vec[n]=lps([m.log_Li,e.log_Xi[n+1],-logz])
    end

    for (n,mrec) in enumerate(e.models)
        m=deserialize(mrec.path)
        p=length(e.posterior_samples)
        θ_mat[:,n+p]=m.θ
        weight_vec[n+p]=lps([m.log_Li, e.log_Xi[end], -log(length(e.models)),-logz])
    end

    kdes=Vector{AbstractKDE}()
    for t in 1:size(θ_mat,1)
        push!(kdes,kde(θ_mat[t,:],weights=exp.(weight_vec)))
    end

    for bv in bivar
        push!(kdes,kde((θ_mat[bv[1],:],θ_mat[bv[2],:]),weights=exp.(weight_vec)))
    end

    return kdes
end


text:
"""
    module SparseUtils.NZVal

For sparse matrices, algorithms that depend only on nzval, m, and n.

This code is generic, for types in `Core`.

The symbols in `NZVal` should *not* be exported, not even for
convenience. They conflict with functions in `Base`.
"""
module NZVal

## This module contains code that depends only on a container
## of non-zero values and the dimensions of the matrix. The
## code references only builtin types.
##
## There is some code in sparsematrix.jl that depends only on the
## number of rows and columns, the collection of nonzero values,
## and possibly the number of nonzero values.
## In particular, the code does not depend on the column and row
## information in colptr and rowval.
## This code is factored out here. It can be used directly in
## CSC, CSR, COO, and possibly other sparse matrix representations.

##  _mapreducezeros in sparsematrix.jl does not depend on any sparse type.
##  So, we import `_mapreducezeros` rather than reproducing it here.
import SparseArrays: _mapreducezeros

@inline numzeros(m, n, numnz) = m * n - numnz

count(pred, m::Integer, n::Integer, nzval::AbstractArray{T}, numnz::Integer) where T =
      Base.count(pred, nzval) +  pred(zero(T)) * numzeros(m, n, numnz)

# sparsematrix.jl: function Base._mapreduce(f, op, ::Base.IndexCartesian, A::SparseMatrixCSC{T}) where T
function _mapreduce(f, op, ::Base.IndexCartesian, m, n, nzval::AbstractArray, numnz)
    return _mapreduce(f, op, n, n, nzval, numnz)
end

function _mapreduce(f, op, m, n, nzval::AbstractArray{T}, numnz) where T
    z = numnz
    nzeros = numzeros(m, n, numnz)
    nlength = m * n
    if z == 0
        if nlength == 0
            Base.mapreduce_empty(f, op, T)
        else
            _mapreducezeros(f, op, T, nzeros - 1, f(zero(T)))
        end
    else
        _mapreducezeros(f, op, T, nzeros, Base._mapreduce(f, op, nzval)) # was nzvalview(A)
    end
end

# sparsematrix.jl: function Base._mapreduce(f, op::typeof(*), A::SparseMatrixCSC{T}) where T
function _mapreduce(f, op::Union{typeof(*), typeof(Base.mul_prod)},
                      m, n, nzval::AbstractArray, numnz=length(nzval))
    nzeros = numzeros(m, n, numnz)
    if nzeros == 0
        # No zeros, so don't compute f(0) since it might throw
        Base._mapreduce(f, op, nzval)
    else
        T = eltype(nzval)
        v = f(zero(T))^(nzeros)
        # Bail out early if initial reduction value is zero
        v == zero(T) ? v : v * Base._mapreduce(f, op, nzval)
    end
end

end # module


text:
using DrWatson
import JSON
using Random
using FileIO
using ArgParse
using Dates

# my_args = Dict(
#     "script" => "simple_network.jl",
#     "algorithm" => "directmc",
#     "directmc_samples" => 2 .^ [3, 4, 5, 6, 7, 8, 9, 10, 11],
#     "run_name" => "2021-06-17",
#     "duration" => 20,
#     "num_responses" => 10_000,
#     "mean_s" => [50],
#     "corr_time_s" => 1,
#     "corr_time_ratio" => 10,
# )

# my_args = Dict(
#     "script" => "simple_network.jl",
#     "algorithm" => "smc",
#     "smc_samples" => 2 .^ [3, 4, 5, 6, 7, 8, 9, 10, 11],
#     "run_name" => "2021-06-21",
#     "duration" => 20,
#     "num_responses" => 10_000,
#     "mean_s" => 50,
#     "corr_time_s" => 1,
#     "corr_time_ratio" => 10,
# )

# my_args = Dict(
#     "script" => "simple_network.jl",
#     "algorithm" => "ti",
#     "ti_samples" => 256,
#     "run_name" => "2021-06-18",
#     "duration" => [1.0, 2.5, 5.0, 7.5, 10.0, 12.5, 15.0, 17.5, 20.0],
#     "num_responses" => 10_000,
#     "mean_s" => 50,
#     "corr_time_s" => 1,
#     "corr_time_ratio" => 10,
# )

# my_args = Dict(
#     "script" => "chemotaxis.jl",
#     "run_name" => "2021-04-05",
#     "smc_samples" => 64,
#     "num_responses" => 5_000,
#     "duration" => 20,
#     "mean_L" => 50,
#     "num_receptors" => 10_000,
#     "Y_tot" => 5000,
#     "LR_timescale" => 0.01,
#     "Y_timescale" => 0.1
# )

my_args = Dict(
    "script" => "cooperative_chemotaxis.jl",
    "algorithm" => "smc",
    "run_name" => "2021-10-18",
    "duration" => 200,
    "tau_l" => [0.02, 0.2, 1.0, 2.0, 5.0, 10.0],
    "num_responses" => 3600,
    "smc_samples" => 128,
)

const NCPUS = 36 * 5
const QUEUE = "highcore"
const NAME = "CHEM_BIG"

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table! s begin
        "--dry-run"
            help = "only show qsub commands"
            action = :store_true
        "--time"
            help = "only show estimated runtime"
            action = :store_true
        "--hold"
            help = "apply user hold at submission"
            action = :store_true
    end
    parse_args(ARGS, s)
end

function DrWatson._wsave(filename, d::Dict)
    if endswith(filename, ".json")
        open(filename, "w") do io
            JSON.print(io, d)
        end
    else
        FileIO.save(filename, d)
    end
end

function submit_job(out_dir, filename, runtime; job_before = nothing, dry_run=false, hold=false)
    jobscript = """
        export JULIA_PROJECT=$(projectdir())

        julia -e "using InteractiveUtils; versioninfo(verbose=true)"
        # julia -J$(projectdir())/GMcmcSysimage.so $(projectdir("scripts", "run_cluster.jl")) $(filename)  $(my_args["script"])
        julia $(projectdir("scripts", "run_cluster.jl")) $(filename)  $(my_args["script"])
        """

    name = NAME
    resources = `-l walltime=$runtime -l select=$NCPUS:ncpus=1:mem=6gb -l place=free`

    if job_before !== nothing
        dependency = `-W depend=afterok:$job_before`
    else
        dependency = ``
    end

    extra_args = hold ? `-h` : ``

    cmd = `qsub -q $QUEUE -N $name $resources $extra_args $dependency -o $out_dir -e $out_dir`

    if dry_run
        println(cmd)
        return "$(rand(1:1000)).head.hollandia.amolf.nl"
    end

    result = ""
    open(cmd, "r+") do io
        print(io, jobscript)
        close(io.in)
        result *= read(io, String)
    end

    print(result)
    strip(result)
end

function estimate_runtime(dict)
    if dict["script"] == "cooperative_chemotaxis.jl"
        factor = 2.0 # gives us some head room

        # we take roughly one hour per sample for the chemotaxis system at SMC M=128
        if dict["algorithm"] == "smc"
            factor *= dict["smc_samples"] / 128
        end

        cput_per_sample = 3800 # very rough

        factor *= dict["duration"] / 200

        return round(Int, factor * cput_per_sample * dict["num_responses"] / NCPUS)
    end

    if dict["script"] == "simple_network.jl"
        if dict["algorithm"] == "annealing"
            factor = 0.14 * 1.5 # empirical factor from AMOLF cluster. The 1.5 is to make sure adequate headroom
        elseif dict["algorithm"] == "ti"
            factor = 0.05 * 2.0 # empirical factor from AMOLF cluster. The 2.0 is to make sure adequate headroom
        elseif dict["algorithm"] == "directmc"
            factor = 0.03 * 2.0 # empirical factor from AMOLF cluster. The 2.0 is to make sure adequate headroom
        elseif dict["algorithm"] == "smc"
            factor = 0.1 * 2
        else
            error("unknown algorithm $(dict["algorithm"])")
        end
        constant = 20 * 60 + NCPUS * 5 # just make sure we have an extra buffer

        num_reactions = dict["duration"] * dict["mean_s"] * (1 / dict["corr_time_s"]) * (1 + dict["corr_time_ratio"])

        val = factor * (num_reactions + 10) * dict["num_responses"]
        return round(Int, val / NCPUS + constant)
    end

    # be very conservative by default
    return 6 * 24 * 60 * 60
end

function submit_sims(; job_before=nothing, dry_run=false, hold=false)
    dicts = dict_list(my_args)

    out_dir = projectdir("data", "output")
    if !dry_run
        filenames = tmpsave(dicts, suffix="json")
        mkpath(out_dir)
    else
        filenames = [randstring() for i=1:length(dicts)]
    end

    for (d, f) in zip(dicts, filenames)
        runtime = estimate_runtime(d)
        job_before = submit_job(out_dir, f, runtime, job_before=job_before, dry_run=dry_run, hold=hold)
    end
end

function print_times()
    sum = Dates.CompoundPeriod()
    dicts = dict_list(my_args)
    for (i, d) in enumerate(dicts)
        runtime = Dates.CompoundPeriod(Second(estimate_runtime(d)))
        println("job $i: $(canonicalize(runtime))")
        sum += runtime
    end

    println("sum: $(canonicalize(sum))")
end

function main()
    parsed_args = parse_commandline()

    if parsed_args["time"]
        print_times()
    else
        submit_sims(dry_run = parsed_args["dry-run"], hold=parsed_args["hold"])
    end
end

main()


text:
export @testmodule

macro testmodule(name::Symbol, ex::Expr)
  const body = Expr(:block,
                    :(eval(x) = Core.eval($name, x)),
                    :(eval(m,x) = Core.eval(m, x)),
                    :(using RegressionTest, RunTests),
                    :(import Base.Test.@test),
                    ex.args...)
  const make_module = Expr(:module, true, esc(name), esc(body))
  const post_process = :(collect_module($(esc(name))))
  return Expr(:toplevel, make_module, post_process)
end


text:
# Contains functions related to distribution networks but not specific to a particular model

## Lookup functions, to build the constraint selection logic

"Return whether the `f_bus` of branch `i` is the reference bus."
function is_frb_branch(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    return haskey(_PM.ref(pm, nw, :frb_branch), i)
end

"Return whether the `f_bus` of ne_branch `i` is the reference bus."
function is_frb_ne_branch(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    return haskey(_PM.ref(pm, nw, :frb_ne_branch), i)
end

"Return whether branch `i` is an OLTC."
function is_oltc_branch(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    return haskey(_PM.ref(pm, nw, :oltc_branch), i)
end

"Return whether ne_branch `i` is an OLTC."
function is_oltc_ne_branch(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    return haskey(_PM.ref(pm, nw, :oltc_ne_branch), i)
end

## Variables

function variable_oltc_branch_transform(pm::_PM.AbstractWModels; kwargs...)
    variable_oltc_branch_transform_magnitude_sqr_inv(pm; kwargs...)
end

"variable: `0 <= ttmi[l]` for `l` in `oltc_branch`es"
function variable_oltc_branch_transform_magnitude_sqr_inv(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    ttmi = _PM.var(pm, nw)[:ttmi] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :oltc_branch)], base_name="$(nw)_ttmi",
        lower_bound = 0.0,
        start = 1.0 / _PM.ref(pm,nw,:oltc_branch,i,"tap")^2
    )

    if bounded
        for (i, br) in _PM.ref(pm, nw, :oltc_branch)
            JuMP.set_lower_bound(ttmi[i], 1.0 / br["tm_max"]^2 )
            JuMP.set_upper_bound(ttmi[i], 1.0 / br["tm_min"]^2 )
        end
    end

    report && _PM.sol_component_value(pm, nw, :branch, :ttmi, _PM.ids(pm, nw, :oltc_branch), ttmi)
end

function variable_oltc_ne_branch_transform(pm::_PM.AbstractWModels; kwargs...)
    variable_oltc_ne_branch_transform_magnitude_sqr_inv(pm; kwargs...)
end

"variable: `0 <= ttmi_ne[l]` for `l` in `oltc_ne_branch`es"
function variable_oltc_ne_branch_transform_magnitude_sqr_inv(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    ttmi_ne = _PM.var(pm, nw)[:ttmi_ne] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :oltc_ne_branch)], base_name="$(nw)_ttmi_ne",
        lower_bound = 0.0,
        start = 1.0 / _PM.ref(pm,nw,:oltc_ne_branch,i,"tap")^2
    )

    if bounded
        for (i, br) in _PM.ref(pm, nw, :oltc_ne_branch)
            JuMP.set_lower_bound(ttmi_ne[i], 1.0 / br["tm_max"]^2 )
            JuMP.set_upper_bound(ttmi_ne[i], 1.0 / br["tm_min"]^2 )
        end
    end

    report && _PM.sol_component_value(pm, nw, :ne_branch, :ttmi, _PM.ids(pm, nw, :oltc_ne_branch), ttmi_ne)
end

## Constraint templates that group several other constraint templates, provided for convenience

function constraint_dist_branch_tnep(pm::_PM.AbstractBFModel, i::Int; nw::Int=_PM.nw_id_default)
    if isempty(ne_branch_ids(pm, i; nw = nw))
        if is_frb_branch(pm, i; nw = nw)
            if is_oltc_branch(pm, i; nw = nw)
                constraint_power_losses_oltc(pm, i; nw = nw)
                constraint_voltage_magnitude_difference_oltc(pm, i; nw = nw)
            else
                constraint_power_losses_frb(pm, i; nw = nw)
                constraint_voltage_magnitude_difference_frb(pm, i; nw = nw)
            end
        else
            _PM.constraint_power_losses(pm, i; nw = nw)
            _PM.constraint_voltage_magnitude_difference(pm, i; nw = nw)
        end
        _PM.constraint_voltage_angle_difference(pm, i; nw = nw)
        _PM.constraint_thermal_limit_from(pm, i; nw = nw)
        _PM.constraint_thermal_limit_to(pm, i; nw = nw)
    else
        expression_branch_indicator(pm, i; nw = nw)
        # constraint_branch_complementarity(pm, i; nw = nw) is best added once per year; if added here, redundant constraints would be generated
        if is_frb_branch(pm, i; nw = nw)
            if is_oltc_branch(pm, i; nw = nw)
                constraint_power_losses_oltc_on_off(pm, i; nw = nw)
                constraint_voltage_magnitude_difference_oltc_on_off(pm, i; nw = nw)
            else
                constraint_power_losses_frb_on_off(pm, i; nw = nw)
                constraint_voltage_magnitude_difference_frb_on_off(pm, i; nw = nw)
            end
        else
            constraint_power_losses_on_off(pm, i; nw = nw)
            constraint_voltage_magnitude_difference_on_off(pm, i; nw = nw)
        end
        _PM.constraint_voltage_angle_difference_on_off(pm, i; nw = nw)
        _PM.constraint_thermal_limit_from_on_off(pm, i; nw = nw)
        _PM.constraint_thermal_limit_to_on_off(pm, i; nw = nw)
    end
end

function constraint_dist_ne_branch_tnep(pm::_PM.AbstractBFModel, i::Int; nw::Int=_PM.nw_id_default)
    if ne_branch_replace(pm, i, nw = nw)
        if is_frb_ne_branch(pm, i, nw = nw)
            if is_oltc_ne_branch(pm, i, nw = nw)
                constraint_ne_power_losses_oltc(pm, i, nw = nw)
                constraint_ne_voltage_magnitude_difference_oltc(pm, i, nw = nw)
            else
                constraint_ne_power_losses_frb(pm, i, nw = nw)
                constraint_ne_voltage_magnitude_difference_frb(pm, i, nw = nw)
            end
        else
            constraint_ne_power_losses(pm, i, nw = nw)
            constraint_ne_voltage_magnitude_difference(pm, i, nw = nw)
        end
        _PM.constraint_ne_thermal_limit_from(pm, i, nw = nw)
        _PM.constraint_ne_thermal_limit_to(pm, i, nw = nw)
    else
        if is_frb_ne_branch(pm, i, nw = nw)
            if is_oltc_ne_branch(pm, i, nw = nw)
                Memento.error(_LOGGER, "addition of a candidate OLTC in parallel to an existing OLTC is not supported")
            else
                constraint_ne_power_losses_frb_parallel(pm, i, nw = nw)
                constraint_ne_voltage_magnitude_difference_frb_parallel(pm, i, nw = nw)
            end
        else
            constraint_ne_power_losses_parallel(pm, i, nw = nw)
            constraint_ne_voltage_magnitude_difference_parallel(pm, i, nw = nw)
        end
        constraint_ne_thermal_limit_from_parallel(pm, i, nw = nw)
        constraint_ne_thermal_limit_to_parallel(pm, i, nw = nw)
    end
    _PM.constraint_ne_voltage_angle_difference(pm, i, nw = nw)
end

## Constraint templates

"Defines voltage drop over a a branch whose `f_bus` is the reference bus"
function constraint_voltage_magnitude_difference_frb(pm::_PM.AbstractBFModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    b_sh_fr = branch["b_fr"]
    tm = branch["tap"]

    constraint_voltage_magnitude_difference_frb(pm, nw, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, tm)
end

"Defines voltage drop over a transformer branch that has an OLTC"
function constraint_voltage_magnitude_difference_oltc(pm::_PM.AbstractBFModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    b_sh_fr = branch["b_fr"]

    constraint_voltage_magnitude_difference_oltc(pm, nw, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr)
end

"Defines branch flow model power flow equations for a branch whose `f_bus` is the reference bus"
function constraint_power_losses_frb(pm::_PM.AbstractBFModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    tm = branch["tap"]
    g_sh_fr = branch["g_fr"]
    g_sh_to = branch["g_to"]
    b_sh_fr = branch["b_fr"]
    b_sh_to = branch["b_to"]

    constraint_power_losses_frb(pm, nw, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, g_sh_to, b_sh_fr, b_sh_to, tm)
end

"Defines branch flow model power flow equations for a transformer branch that has an OLTC"
function constraint_power_losses_oltc(pm::_PM.AbstractBFModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    g_sh_to = branch["g_to"]
    b_sh_fr = branch["b_fr"]
    b_sh_to = branch["b_to"]

    constraint_power_losses_oltc(pm, nw, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, g_sh_to, b_sh_fr, b_sh_to)
end

""
function constraint_ne_power_losses_frb(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :ne_branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    tm = branch["tap"]
    g_sh_fr = branch["g_fr"]
    g_sh_to = branch["g_to"]
    b_sh_fr = branch["b_fr"]
    b_sh_to = branch["b_to"]

    vad_min = _PM.ref(pm, nw, :off_angmin)
    vad_max = _PM.ref(pm, nw, :off_angmax)

    constraint_ne_power_losses_frb(pm, nw, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, g_sh_to, b_sh_fr, b_sh_to, tm, vad_min, vad_max)
end

""
function constraint_ne_power_losses_oltc(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :ne_branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    g_sh_to = branch["g_to"]
    b_sh_fr = branch["b_fr"]
    b_sh_to = branch["b_to"]

    vad_min = _PM.ref(pm, nw, :off_angmin)
    vad_max = _PM.ref(pm, nw, :off_angmax)

    constraint_ne_power_losses_oltc(pm, nw, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, g_sh_to, b_sh_fr, b_sh_to, vad_min, vad_max)
end

"""
Defines voltage drop over a branch, linking from and to side voltage magnitude
"""
function constraint_ne_voltage_magnitude_difference_frb(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :ne_branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    b_sh_fr = branch["b_fr"]
    tm = branch["tap"]

    constraint_ne_voltage_magnitude_difference_frb(pm, nw, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, tm)
end

"""
Defines voltage drop over a branch, linking from and to side voltage magnitude
"""
function constraint_ne_voltage_magnitude_difference_oltc(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :ne_branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    b_sh_fr = branch["b_fr"]
    tm_min = branch["tm_min"]
    tm_max = branch["tm_max"]

    constraint_ne_voltage_magnitude_difference_oltc(pm, nw, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, tm_min, tm_max)
end

## Constraint implementations not limited to a specific model type

""
function constraint_ne_voltage_magnitude_difference_frb(pm::_PM.AbstractBFAModel, n::Int, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, tm)
    branch = _PM.ref(pm, n, :ne_branch, i)
    fr_bus = _PM.ref(pm, n, :bus, f_bus)
    to_bus = _PM.ref(pm, n, :bus, t_bus)
    M_hi   =  1.0^2/tm^2 - to_bus["vmin"]^2
    M_lo   = -1.0^2/tm^2 + to_bus["vmax"]^2
    p_fr = _PM.var(pm, n, :p_ne, f_idx)
    q_fr = _PM.var(pm, n, :q_ne, f_idx)
    w_to = _PM.var(pm, n, :w, t_bus)
    z    = _PM.var(pm, n, :branch_ne, i)
    # w_fr is assumed equal to 1.0

    JuMP.@constraint(pm.model, (1.0/tm^2) - w_to <= 2*(r*p_fr + x*q_fr) + M_hi*(1-z) )
    JuMP.@constraint(pm.model, (1.0/tm^2) - w_to >= 2*(r*p_fr + x*q_fr) - M_lo*(1-z) )
end

""
function constraint_ne_voltage_magnitude_difference_oltc(pm::_PM.AbstractBFAModel, n::Int, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, tm_min, tm_max)
    branch = _PM.ref(pm, n, :ne_branch, i)
    fr_bus = _PM.ref(pm, n, :bus, f_bus)
    to_bus = _PM.ref(pm, n, :bus, t_bus)
    M_hi   =  1.0^2/tm_min^2 - to_bus["vmin"]^2
    M_lo   = -1.0^2/tm_max^2 + to_bus["vmax"]^2
    p_fr = _PM.var(pm, n, :p_ne, f_idx)
    q_fr = _PM.var(pm, n, :q_ne, f_idx)
    ttmi = _PM.var(pm, n, :ttmi_ne, i)
    w_to = _PM.var(pm, n, :w, t_bus)
    z    = _PM.var(pm, n, :branch_ne, i)
    # w_fr is assumed equal to 1.0 to preserve the linearity of the model

    JuMP.@constraint(pm.model, 1.0*ttmi - w_to <= 2*(r*p_fr + x*q_fr) + M_hi*(1-z) )
    JuMP.@constraint(pm.model, 1.0*ttmi - w_to >= 2*(r*p_fr + x*q_fr) - M_lo*(1-z) )
end


text:
using Test

include("sudokutest.jl")


text:
function allocate(::Type{oneL0.DeviceBuffer}, ctx, dev, bytes::Int, alignment::Int)
    bytes == 0 && return oneL0.DeviceBuffer(ZE_NULL, bytes, ctx, dev)

    buf = device_alloc(ctx, dev, bytes, alignment)
    make_resident(ctx, dev, buf)

    return buf
end

function release(buf::oneL0.AbstractBuffer)
    sizeof(buf) == 0 && return

    ctx = oneL0.context(buf)
    dev = oneL0.device(buf)

    evict(ctx, dev, buf)
    free(buf; policy=oneL0.ZE_DRIVER_MEMORY_FREE_POLICY_EXT_FLAG_BLOCKING_FREE)

    # TODO: queue-ordered free from non-finalizer tasks once we have
    #       `zeMemFreeAsync(ptr, queue)`

    return
end


text:
using ProgressMeter
using Feather
include(joinpath("src", "module.jl"))

begin
    BATCH_NAME = "original"
    GRID_DIMS = (10, 10)
    STUBBORN_COUNT = 4
    STEPS = 100
    REPLICATES = 2
    WHEN = 10
    WRITE = true
end

begin
    baseline = NTuple{2, Int}[]
    line_edge = [(1, 1), (1, 2), (1, 3), (1, 4)]
    line_center = [(5, 3), (5, 4), (5, 5), (5, 6)]
    square_corner = [(1, 1), (1, 2), (2, 1), (2, 2)]
    square_center = [(5, 5), (5, 6), (6, 5), (6, 6)]
    corners = [(1, 1), (1, 10), (10, 1), (10, 10)]
    diagonal = [(1, 1), (4, 4), (7, 7), (10, 10)]
    distance_center = [(4, 4), (4, 7), (7, 4), (7, 7)]
    config_list = [
        baseline, line_edge, line_center, square_corner,
        square_center, corners, diagonal, distance_center
    ]
    filename_list = [
        "baseline", "line_edge", "line_center", "square_corner",
        "square_center", "corners", "diagonal", "distance_center"
    ]
    AxelFraud.create_config_table(
        config_list=config_list,
        grid_dims=GRID_DIMS,
        batch_name=BATCH_NAME,
        write=WRITE
    )
end

begin
    p = ProgressMeter.Progress(length(config_list), "Running Simulations ...")
    for (config, filename) in zip(config_list, filename_list)
        agent_df = AxelFraud.run_config(
            config_name=filename,
            batch_name=BATCH_NAME,
            grid_dims=GRID_DIMS,
            stubborn_positions=config,
            steps=STEPS,
            replicates=REPLICATES,
            when=WHEN,
            write=WRITE
        )
        ProgressMeter.next!(p)
    end
end

random_agent_df, random_configs = AxelFraud.run_random(
    config_name="random",
    batch_name=BATCH_NAME,
    grid_dims=GRID_DIMS,
    stubborn_count=STUBBORN_COUNT,
    steps=STEPS,
    replicates=REPLICATES,
    when=WHEN,
    rnd_seed=0,
    write=false
)
Feather.write(joinpath("data", BATCH_NAME, "random_configs.feather"), random_configs)


text:
# using Unitful
# import Unitful: ML, ha, mm

# @unit ML/ha "ML/ha" Megaliter/hectare (1//100)ha false

const ML_to_mm = 100.0
const mm_to_ML = 100.0

max_stars_repo_name: ConnectedSystems/agtor-jl

max_issues_repo_name: ConnectedSystems/agtor-jl

max_forks_repo_name: ConnectedSystems/agtor-jl


text:
# SPDX-License-Identifier: X11
# 2020-11-05
# Trapezoid Sum (200pt)

sumto(n) = n * (n + 1) ÷ 2
sumfromto(n, m) = sumto(m) - sumto(n) + n

function main()
  N = parse(Int, readline())
  sum = 0
  for i ∈ 1:N
    A,B = parse.(Int, split(readline()))
    sum += sumfromto(A,B)
  end
  println(sum)
end

main()

max_stars_repo_name: jthvai/competitive-programming

max_issues_repo_name: jthvai/competitive-programming

max_forks_repo_name: jthvai/competitive-programming


text:
"""
    classificationhead(n, inchannels)

Head for convolutional networks with `inchannels` output channels.
Returns vector of size `n` with `softmax` applied.
"""
function classificationhead(n, inchannels)
    return Chain(
        AdaptiveMeanPool((1,1)),
        flatten,
        Dense(inchannels, n),
    )
end


text:
# using Test
# using StaticArrays
# using TrajectoryOptimization
# using RobotZoo: Cartpole
# using ForwardDiff
# using BenchmarkTools
# using RobotDynamics
# using DiffResults
# using FiniteDiff
# using LinearAlgebra
import TrajectoryOptimization: stage_cost, CostFunction, Expansion
# const TO = TrajectoryOptimization

##
struct CartpoleCost{T} <: TO.CostFunction
    Q::Vector{T}
    R::Vector{T}
end
function TO.stage_cost(cost::CartpoleCost, x, u)
    return TO.stage_cost(cost, x) + 0.5 * cost.R[1] * u[1]^2
end
function TO.stage_cost(cost::CartpoleCost, x)
    y = x[1]
    θ = x[2]
    ydot = x[3]
    θdot = x[4]
    J = cost.Q[2] * cos(θ/2)
    J += 0.5* (cost.Q[1] * y^2 + cost.Q[3] * ydot^2 + cost.Q[4] * θdot^2)
    return J
end
TO.state_dim(::CartpoleCost) = 4
TO.control_dim(::CartpoleCost) = 1

## Initialize Model
@testset "Nonlinear Costs" begin
    model = Cartpole()
    n,m = size(model)
    x = zeros(n)
    u = zeros(m)
    cst = CartpoleCost([1,2,3,4.], [2.])
    @test TO.stage_cost(cst, x, u) ≈ 2.0
    x[2] = π
    @test TO.stage_cost(cst, x) ≈ 0 atol = 10*eps()
    x,u = rand(model)
    z = KnotPoint(x,u,0.1)
    zterm = KnotPoint(x,u,0.0)

    ## Test ForwardDiff
    TO.diffmethod(::CartpoleCost) = RobotDynamics.ForwardAD()
    E = Expansion{Float64}(n,m)
    TO.gradient!(E, cst, zterm)
    TO.hessian!(E, cst, zterm)
    TO.gradient!(E, cst, z)
    TO.hessian!(E, cst, z)
    @test (@allocated TO.gradient!(E, cst, zterm)) == 0
    @test (@allocated TO.hessian!(E, cst, zterm)) == 0
    @test (@allocated TO.gradient!(E, cst, z)) == 0
    @test (@allocated TO.hessian!(E, cst, z)) == 0
    @test E.xx ≈ Diagonal([1, -0.5*cos(x[2]/2), 3, 4])
    @test E.uu ≈ Diagonal([2])
    @test E.x ≈ [1*x[1], -sin(x[2]/2), 3*x[3], 4*x[4]]
    @test E.u ≈ 2*u

    TO.diffmethod(::CartpoleCost) = RobotDynamics.FiniteDifference()
    E2 = Expansion{Float64}(n,m)
    TO.gradient!(E2, cst, zterm)
    TO.hessian!(E2, cst, zterm)
    TO.gradient!(E2, cst, z)
    TO.hessian!(E2, cst, z)

    cache = TO.ExpansionCache(cst)
    TO.gradient!(E2, cst, zterm, cache)
    TO.hessian!(E2, cst, zterm, cache)
    TO.gradient!(E2, cst, z, cache)
    TO.hessian!(E2, cst, z, cache)

    @test (@allocated TO.gradient!(E2, cst, zterm, cache)) == 0
    @test (@allocated TO.hessian!(E2, cst, zterm, cache)) == 0
    @test (@allocated TO.gradient!(E2, cst, z, cache)) == 0
    @test (@allocated TO.hessian!(E2, cst, z, cache)) == 0

    @test norm(E2.hess- E.hess) < 1e-6
    @test norm(E2.grad- E.grad) < 1e-6

    ## Test expansion of trajectory
    N = 11
    obj = Objective(cst, N)
    X = [rand(n) for k = 1:N]
    U = [rand(m) for k = 1:N-1]
    Z = Traj(X,U,fill(0.1,N))

    TO.diffmethod(::CartpoleCost) = RobotDynamics.ForwardAD()
    E0 = [Expansion{Float64}(n,m) for k = 1:N]
    cache = TO.ExpansionCache(obj)
    TO.cost_expansion!(E0, obj, Z, cache, init=true, rezero=true)
    # @test (@allocated TO.cost_expansion!(E0, obj, Z, cache, init=true, rezero=true)) == 0

    TO.diffmethod(::CartpoleCost) = RobotDynamics.FiniteDifference()
    E1 = [Expansion{Float64}(n,m) for k = 1:N]
    cache = TO.ExpansionCache(obj)
    TO.cost_expansion!(E1, obj, Z, cache, init=true, rezero=true)
    # @test (@allocated TO.cost_expansion!(E1, obj, Z, cache, init=true, rezero=true)) == 0
end


text:
# Needed for defining the rule for `adjoint(A::ITensor)`
# which currently doesn't work by overloading `ChainRulesCore.rrule`
using ZygoteRules: @adjoint

@adjoint function Base.adjoint(x::ITensor)
  y = prime(x)
  function setinds_pullback(ȳ)
    x̄ = ITensors.setinds(ȳ, inds(x))
    return (x̄,)
  end
  return y, setinds_pullback
end

max_stars_repo_name: mtfishman/ITensorsGrad

max_issues_repo_name: mtfishman/ITensorsGrad

max_forks_repo_name: mtfishman/ITensorsGrad


text:
function preproc_atari_frame(s)
    r1 = 36:2:194
    r2 = 1:2:160
    v = (s[r1, r2, 1] .+  s[r1, r2, 2] .+ s[r1, r2, 3]) .÷ UInt8(3)
    reshape(v, size(v)..., 1, 1)
end

AtariPOMDP(environment; version = :v0, frame_stack = 4) = GymPOMDP(environment, version = version, frame_stack = frame_stack, pixel_observations = true, special_render = preproc_atari_frame, sign_reward = true)

max_forks_repo_name: ancorso/POMDPGym


text:
"""
Helps to maintain the pool of reusable arrays of different sizes
and reduce the burden on garbage collection.
"""
immutable ArrayPool{T<:Number}
    length_pools::Dict{Int, Vector{Vector{T}}}

    ArrayPool{T}() where {T<:Number} = new(Dict{Int, Vector{Vector{T}}}())
end

"""
Gets an array of specific size from the pool.
The returned array should be returned back to the pool using `release!()`.
"""
function acquire!{T}(pool::ArrayPool{T}, size)
    len = prod(size)
    len_pool = haskey(pool.length_pools, len) ?
               pool.length_pools[len] :
               get!(pool.length_pools, len, Vector{Vector{T}}())
    isempty(len_pool) ? Array{T}(size) : reinterpret(T, pop!(len_pool), size)
end

"""
Releases an array returned by `acquire!()` back into the pool.
"""
function release!{T}(pool::ArrayPool{T}, arr::Array{T})
    len = length(arr)
    len_pool = haskey(pool.length_pools, len) ?
               pool.length_pools[len] :
               get!(pool.length_pools, len, Vector{Vector{T}}())
    push!(len_pool, reinterpret(T, arr, (len,)))
    pool
end


text:
using Test
using JSON3
using Pinecone

APIKEY = ENV["PINECONE_API_KEY"]
context = Pinecone.init(APIKEY, "us-west1-gcp")

TESTINDEX = "filter-example"
NAMESPACE = "mynamespace"

testdict = Dict{String, Any}("genre"=>"documentary", "year"=>2019);
v1 = [0.3,0.11,0.3,0.3,0.3,0.3,0.3,0.3,0.4,0.3]
v2 = [0.9, 0.8, 0.7, 0.6, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3]

@testset verbose = true "Test query()" begin
    testvector = Pinecone.PineconeVector("testid", v1, testdict)
    index = Pinecone.Index(TESTINDEX)
    #meta = [Dict{String,Any}("foo"=>"bar"), Dict{String,Any}("bar"=>"baz")]
    testvector = Pinecone.PineconeVector("testid", v1, testdict)
    testvector2 = Pinecone.PineconeVector("testid2", v2, testdict)

    result = Pinecone.upsert(context, index, [testvector, testvector2], "")
    result = Pinecone.upsert(context, index, [testvector, testvector2], NAMESPACE)
    #basic test for results
    result = Pinecone.query(context, index, [testvector], 4)
    @test result !== nothing
    @test typeof(result) == String
    rvjson = JSON3.read(result)
    @test haskey(rvjson, "results")
    @test length(rvjson["results"]) > 0

    #basic test for results with a namespace
    result = Pinecone.query(context, index, [testvector], 4, NAMESPACE)
    @test result !== nothing
    @test typeof(result) == String
    rvjson = JSON3.read(result)
    @test haskey(rvjson, "results")
    @test length(rvjson["results"]) > 0

    #test asking for values
    result = Pinecone.query(context, index, [v1, v1], 4, "", true)
    @test result !== nothing
    @test typeof(result) == String
    rvjson = JSON3.read(result)
    @test haskey(rvjson, "results")
    @test length(rvjson["results"]) > 0
    @test haskey(rvjson["results"][1]["matches"][1], "values")
    @test length(rvjson["results"][1]["matches"][1]["values"]) > 0

      #test asking for values but with namespace
      result = Pinecone.query(context, index, [v1, v1], 4, NAMESPACE, true)
      @test result !== nothing
      @test typeof(result) == String
      rvjson = JSON3.read(result)
      @test haskey(rvjson, "results")
      @test length(rvjson["results"]) > 0
      @test haskey(rvjson["results"][1]["matches"][1], "values")
      @test length(rvjson["results"][1]["matches"][1]["values"]) > 0

    #test no values returned
    result = Pinecone.query(context, index, [v1, v1], 4, "", false)
    @test result !== nothing
    @test typeof(result) == String
    rvjson = JSON3.read(result)
    @test haskey(rvjson, "results")
    @test length(rvjson["results"]) > 0
    @test haskey(rvjson["results"][1]["matches"][1], "values")
    @test length(rvjson["results"][1]["matches"][1]["values"]) == 0

        #test no values returned with namespace
        result = Pinecone.query(context, index, [v1, v1], 4, NAMESPACE, false)
        @test result !== nothing
        @test typeof(result) == String
        rvjson = JSON3.read(result)
        @test haskey(rvjson, "results")
        @test length(rvjson["results"]) > 0
        @test haskey(rvjson["results"][1]["matches"][1], "values")
        @test length(rvjson["results"][1]["matches"][1]["values"]) == 0

     #test metadata returned but values not
     result = Pinecone.query(context, index, [v1, v1], 4, "", false, true)
     @test result !== nothing
     @test typeof(result) == String
     rvjson = JSON3.read(result)
     @test haskey(rvjson, "results")
     @test length(rvjson["results"]) > 0
     @test haskey(rvjson["results"][1]["matches"][1], "metadata")
     @test length(rvjson["results"][1]["matches"][1]["metadata"]) > 0

          #test metadata returned but values not with namespace
          result = Pinecone.query(context, index, [v1, v1], 4, NAMESPACE, false, true)
          @test result !== nothing
          @test typeof(result) == String
          rvjson = JSON3.read(result)
          @test haskey(rvjson, "results")
          @test length(rvjson["results"]) > 0
          @test haskey(rvjson["results"][1]["matches"][1], "metadata")
          @test length(rvjson["results"][1]["matches"][1]["metadata"]) > 0

    #test metadata returned and values returned
    result = Pinecone.query(context, index, [v1, v1], 4, "", true, true)
    @test result !== nothing
    @test typeof(result) == String
    rvjson = JSON3.read(result)
    @test haskey(rvjson, "results")
    @test length(rvjson["results"]) > 0
    @test haskey(rvjson["results"][1]["matches"][1], "metadata")
    @test length(rvjson["results"][1]["matches"][1]["metadata"]) > 0

    #test metadata returned and values returned with namespace
    result = Pinecone.query(context, index, [v1, v1], 4, NAMESPACE, true, true)
    @test result !== nothing
    @test typeof(result) == String
    rvjson = JSON3.read(result)
    @test haskey(rvjson, "results")
    @test length(rvjson["results"]) > 0
    @test haskey(rvjson["results"][1]["matches"][1], "metadata")
    @test length(rvjson["results"][1]["matches"][1]["metadata"]) > 0

    #test metadata not returned but values returned
    result = Pinecone.query(context, index, [v1, v1], 4, "", true, false)
    @test result !== nothing
    @test typeof(result) == String
    rvjson = JSON3.read(result)
    @test haskey(rvjson, "results")
    @test length(rvjson["results"]) > 0
    @test !haskey(rvjson["results"][1]["matches"][1], "metadata")

      #test metadata not returned but values returned with namespace
      result = Pinecone.query(context, index, [v1, v1], 4, NAMESPACE, true, false)
      @test result !== nothing
      @test typeof(result) == String
      rvjson = JSON3.read(result)
      @test haskey(rvjson, "results")
      @test length(rvjson["results"]) > 0
      @test !haskey(rvjson["results"][1]["matches"][1], "metadata")

    @testset "Test Filters" begin
      #...means the "genre" takes on both values.
      moviemeta = [Dict{String, Any}("genre"=>["comedy","documentary"]), Dict{String, Any}("genre"=>["comedy","documentary"])]
      result = Pinecone.upsert(context, index, ["zipA", "zipB"], [[0.1, 0.2, 0.3, 0.4, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3],
[0.9, 0.8, 0.7, 0.6, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3]], moviemeta, "mynamespace")
      filter = """{
        "genre": {
          "\$in": [
            "comedy",
            "documentary",
            "drama"
          ]
        },
        "year": {
          "\$eq": 2019
        }
      }"""
      result = Pinecone.query(context, index, [v1], 4, "mynamespace", true, true, JSON3.read(filter, Dict{String, Any}))
      rvjson = JSON3.read(result)
      @test haskey(rvjson, "results")
      @test length(rvjson["results"]) > 0
      @test length(rvjson["results"][1]["matches"]) > 0

      #test other version that takes PineconeVector
      result = Pinecone.query(context, index, [testvector], 4, "mynamespace", true, true, JSON3.read(filter, Dict{String, Any}))
      rvjson = JSON3.read(result)
      @test haskey(rvjson, "results")
      @test length(rvjson["results"]) > 0
      @test length(rvjson["results"][1]["matches"]) > 0

      #changing to filter to match nothing, now checking that.  Should get no matches since no 'action' in the moviemeta
      filter = """{
        "genre": {
          "\$in": [
            "action"
          ]
        },
        "year": {
          "\$eq": 2019
        }
      }"""
      result = Pinecone.query(context, index, [v1], 4, "mynamespace", true, true, JSON3.read(filter, Dict{String, Any}))
      rvjson = JSON3.read(result)
      @test haskey(rvjson, "results")
      @test length(rvjson["results"]) > 0
      @test length(rvjson["results"][1]["matches"]) == 0

      result = Pinecone.query(context, index, [testvector], 4, "mynamespace", true, true, JSON3.read(filter, Dict{String, Any}))
      println(JSON3.@pretty(result))
      rvjson = JSON3.read(result)
      @test haskey(rvjson, "results")
      @test length(rvjson["results"]) > 0
      @test length(rvjson["results"][1]["matches"]) == 0
    end

    result = Pinecone.query(context, index, [v1, v1], 4)
    @test result !== nothing
    @test typeof(result) == String
    #should get nothing here w/ this bogus namespace
    result = Pinecone.query(context, index, [v1, v1], 4, "bogusempty", true)
    #the bogus namespace returns HTTP 400 but with JSON body.  query() returns nothing so check
    rvjson = JSON3.read(result)
    resultsarr = rvjson["results"];
    @test length(resultsarr) == 2
    #now returns empty results w/ matches that'll have empty array, make sure bogus query is empty
    @test length(resultsarr[1]["matches"]) == 0

    println("***************** Test query() topK *****************")
    # test topk exceeded
    @test_throws ArgumentError Pinecone.query(context, index, [testvector], 10001)

    println("***************** Test query() topk exceeded include results*************")
    # test topk exceeded when values included in return results
    @test_throws ArgumentError Pinecone.query(context, index, [testvector], 10000, "", true)

    # test topk exceeded with alternate query form
    @test_throws ArgumentError Pinecone.query(context, index, [v1], 10001)
    # test topk exceeded when values included in return results
    @test_throws ArgumentError Pinecone.query(context, index, [v1], 10000, "", true)

    # max top k when including data or metadata is 1000, test for 1001 results
    # test topk exceeded when values included in return results with includesvalues=true
    @test_throws ArgumentError Pinecone.query(context, index, [v1], 1001, "", true, false)
    # test topk exceeded when values included in return results with includesmeta=true
    @test_throws ArgumentError Pinecone.query(context, index, [v1], 1001, "", false, true)

 end


text:
# A simple example to use `BeeEncoder.jl`

using BeeEncoder

@beeint x  0 5
@beeint y -4 9
@beeint z -5 10

@constrain x + y == z

@beeint w 0 10

xl = @beebool x[1:4]

@constrain xl[1] == -xl[2]
@constrain xl[2] == true

@constrain sum([-xl[1], xl[2], -xl[3], xl[4]]) == w

BeeEncoder.render()

max_stars_repo_name: newptcai/BeeEncoder

max_issues_repo_name: newptcai/BeeEncoder

max_forks_repo_name: newptcai/BeeEncoder


text:
using ZfpCompression
using Test

@testset "Lossless 1-4D for all types" begin

    for T in (Float32,Float64)
        # test 1-4D
        for sizes in [(100,),(100,50),(20,30,40),(20,20,20,20)]
            A = rand(T,sizes...)
            Ac = zfp_compress(A)
            Ad = zfp_decompress(Ac)
            @test Ad == A
        end
    end

    for T in (Int32,Int64)
        # test 2-4D
        for sizes in [(100,50),(20,30,40),(20,20,20,20)]
            A = rand(T,sizes...)
            Ac = zfp_compress(A)
            Ad = zfp_decompress(Ac)
            @test Ad == A
        end
    end
end

@testset "Max abs error is bound in 1-4D for floats" begin

    for tol in [1e-1,1e-3,1e-5,1e-7]
        for T in (Float32,Float64)
            # test 1-4D
            for sizes in [(100,),(100,50),(20,30,40),(20,20,20,20)]
                A = rand(T,sizes...)
                Ac = zfp_compress(A;tol)
                Ad = zfp_decompress(Ac)
                @test maximum(abs.(Ad-A)) <= tol
            end
        end
    end
end


text:
# Copyright (c) 2013: Iain Dunning, Miles Lubin, and contributors
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

using Test

function runtests(mod)
    for name in names(mod; all = true)
        if !startswith("$(name)", "test_")
            continue
        end
        @testset "$(name)" begin
            getfield(mod, name)()
        end
    end
end

@testset "C" begin
    include("C_wrapper.jl")
end

@testset "MathOptInterface" begin
    include("MOI_wrapper.jl")
end