import KnetLayers: _batchSizes2indices, _pack_sequence, arrtype
#####
##### Abstract Type Definitions
#####

"""
    Model

    Abstract type for all the models in Morse Package
    see `subtypes(Model)`
"""
abstract type Model;end;

"""
    Sequential <: Model

    Abstract type for all sequential models in Morse
    see `subtypes(Sequential)`
"""
abstract type Sequential <: Model;end;

function predict(M::Model, sentence::String; v::Vocabulary, p::Parser, maxL=20)
    encoded_ios = map(w->EncodedIO(collect(w),v=v,mask=v.specialTokens.unk,len=maxL), split(sentence))
    d = miniBatch(encoded_ios, v)
    preds,_ = predict(M, d; v=v)
    return map(enumerate(d.encodedIOs)) do (i,encodedIO)
        if isa(M,Sequential)
           @inbounds pred  = StringAnalysis(encodedIO.chars, preds[i,:], v=v)
        else
           @inbounds pred  = StringAnalysis(encodedIO.chars, preds[i], v=v)
        end
        makeFormat(pred,p)
    end
end

"""
    Discriminative <: Model

    Abstract type for all discriminative models in Morse
    see `subtypes(Discriminative)`
"""
abstract type Discriminative <: Model;end;

#####
##### Primitives
#####

"""
    WordEncoder(config::Dict,v::Vocabulary)
    (M::WordEncoder)(input::Vector{Int}, batchSizes::Vector{Int},
                     unsortingIndices::Vector{Int}; training=false)

    The `WordEncoder` described in the paper.
"""
struct WordEncoder
    embed::Embed
    dropout::Dropout
    encoder::LSTM
end

WordEncoder(o::Dict,v::Vocabulary) = WordEncoder(
    Embed(input=length(v.chars), output=o[:embedSizes][1]),
    Dropout(p=o[:dropouts]),
    LSTM(input=o[:embedSizes][1], hidden=o[:hiddenSizes][1], seed=o[:seed]) |> fgbias!
)

# function (M::WordEncoder)(inputs::Vector{Int}, sizes::Vector{Int}, indices::Vector{Int}; training=false)
#     out = M.encoder(M.dropout(M.embed(inputs), enable=training), batchSizes=sizes, hy=true, cy=true)
#     return out.hidden[:,:,end][:,indices], out.memory[:,:,end][:,indices]
# end

function (M::WordEncoder)(inputs::Vector{Int}, sizes::Vector{Int}, indices::Vector{Int}; training=false)
    if value(M.encoder.params) isa KnetArray
        out = M.encoder(M.dropout(M.embed(inputs), enable=training), batchSizes=sizes, hy=true, cy=true)
        return out.hidden[:,:,end][:,indices], out.memory[:,:,end][:,indices]
    else
        hs=[]; cs=[];
        for ind in _batchSizes2indices(sizes)
            out = M.encoder(reshape(M.embed(inputs[ind]),size(M.embed.weight,1),1,length(ind)); hy=true, cy=true)
            push!(hs, reshape(out.hidden,size(out.hidden,1),1))
            push!(cs, reshape(out.memory,size(out.memory,1),1))
        end
        return hcat(hs[indices]...), hcat(cs[indices]...)
    end
end

"""
    ContextEncoder(config::Dict,v::Vocabulary)
    (M::ContextEncoder)(hiddens; training=false)

    The `ContextEncoder` described in the paper.
"""
struct ContextEncoder
    encoder::LSTM
    reducer::Dense
    dropout::Dropout
end

ContextEncoder(o::Dict,v::Vocabulary) = ContextEncoder(
    LSTM(input=o[:hiddenSizes][1], hidden=o[:hiddenSizes][2], bidirectional=true, seed=o[:seed]) |> fgbias!,
    Dense(input=2o[:hiddenSizes][2],output=o[:hiddenSizes][3], activation=ReLU()),
    Dropout(p=o[:dropouts])
)

function (M::ContextEncoder)(hiddens; training=false)
    input = reshape(hiddens,size(hiddens,1),1,size(hiddens,2))
    y = M.encoder(M.dropout(input; enable=training)).y
    h = M.reducer(reshape(y,size(y,1),size(y,3)))
    return h, zero(value(h)) #TO-DO: zeros
end

"""
    OutputEncoder(config::Dict,v::Vocabulary)
    (M::OutputEncoder)(tagEmbeddings::Vector)
    (M::OutputEncoder)(emb, h, c)

    The `OutputEncoder` described in the paper.
"""
struct OutputEncoder
    encoder::LSTM
    dropout::Dropout
    previous::Int
    hidden
end

OutputEncoder(o::Dict,v::Vocabulary) = OutputEncoder(
    LSTM(input=o[:embedSizes][2], hidden=o[:hiddenSizes][4], seed=o[:seed]) |> fgbias!,
    Dropout(p=o[:dropouts]),
    o[:previous],
    arrtype(zeros(o[:hiddenSizes][4],1,1))
)

function (M::OutputEncoder)(emb, h, c)
    y = M.encoder(emb, h, c; hy=true, cy=true)
    return y.hidden, y.memory
end

function (M::OutputEncoder)(tagEmbeddings::Vector)
    hos, cos = [], [];
    H,_,_ = size(M.hidden) # assume one layer output encoder
    for i=1:length(tagEmbeddings)
        ho, co = zero(M.hidden), zero(M.hidden) #TO-DO zeros
        for j = max(i-M.previous,1):i-1
             out = M.encoder(M.dropout(tagEmbeddings[j],enable=true), ho, co; hy=true, cy=true)
             ho, co = out.hidden, out.memory
        end
        push!(hos,reshape(ho,H,1))
        push!(cos,reshape(co,H,1))
    end
    return hcat(hos...), hcat(cos...)
end

"""
    Decoder(config::Dict,v::Vocabulary)
    (M::Decoder)(input, hiddens1, hiddens2; training=false)

    The two layer `Decoder` described in the paper.
"""
struct Decoder
    L1::LSTM
    L2::LSTM
    dropout::Dropout
end

function make3D(x)
    if ndims(x) < 3
        return reshape(x,size(x)...,1)
    end
    return x
end

function (M::Decoder)(input, hiddens1, hiddens2; training=false)
    out1   = M.L1(M.dropout(input; enable=training), make3D.(hiddens1)...)
    out2   = M.L2(M.dropout(out1.y; enable=training), make3D.(hiddens2)...)
    return (out1.hidden, out1.memory),
           (M.dropout(out2.hidden; enable=training), out2.memory)
end

Decoder(o::Dict,v::Vocabulary) = Decoder(
    LSTM(input=o[:embedSizes][2], hidden=o[:hiddenSizes][3], seed=o[:seed]) |> fgbias!,
    LSTM(input=o[:hiddenSizes][3], hidden=o[:hiddenSizes][3], seed=o[:seed]) |> fgbias!,
    Dropout(p=o[:dropouts])
)

"""
    recurrentLoss(M::Sequential, d::SentenceBatch, hiddens1, hiddens2; startIndex=4)

    Calculate recurrent negative likelihood for sequential models
"""
function recurrentLoss(M::Sequential, d::SentenceBatch, hiddens1, hiddens2; startIndex=4)
    total = 0.0
    input = fill(startIndex, size(d.seqOutputs,1))
    for i=1:size(d.seqOutputs,2)
        hiddens1, hiddens2 = M.decoder(M.outputEmbed(input), hiddens1, hiddens2; training=true)
        h2 = hiddens2[1]
        scores = M.output(reshape(h2,size(h2,1),size(h2,2)))
        # loss calculation
        total += M.loss(scores, d.seqOutputs[:,i] .* d.masks[:,i]; average=false)
        input  = d.seqOutputs[:,i]
    end
    return total
end

"""
    recurrentPredict(M::Sequential, d::SentenceBatch, hiddens1, hiddens2; startIndex=4)

    Make prediciton with recurrent decoder for sequential models except MorseModel due to its `OutputEncoder`
"""
function recurrentPredict(M::Sequential, d::SentenceBatch, hiddens1, hiddens2; startIndex=4)
    wordNumber, timeSteps = size(d.seqOutputs)
    preds = zeros(Int, wordNumber, timeSteps)
    input = fill(startIndex, wordNumber)
    total = 0.0
    for i=1:timeSteps
        hiddens1, hiddens2 = M.decoder(M.outputEmbed(input), hiddens1, hiddens2; training=false)
        h2 = hiddens2[1]
        scores = M.output(reshape(h2,size(h2,1),size(h2,2)))
        # prediction
        input = vec(getindex.(argmax(convert(Array,scores),dims=1),1))
        preds[:,i] = input
        #loss
        total += M.loss(scores, d.seqOutputs[:,i] .* d.masks[:,i]; average=false)
    end
    return preds, total
end

#####
##### Sequential Models
#####

"""
    MorseModel <: Sequential
    MorseModel(config::Dict,v::Vocabulary)
    loss(M::MorseModel, d::SentenceBatch; v::Vocabulary)
    predict(M::MorseModel, d::SentenceBatch; v::Vocabulary)

    The `MorseModel` described in the paper
"""
struct MorseModel <: Sequential
    wordEncoder::WordEncoder
    contextEncoder::ContextEncoder
    outputEncoder::OutputEncoder
    outputEmbed::Embed
    decoder::Decoder
    output::Linear
    loss::CrossEntropyLoss
end

MorseModel(o::Dict,v::Vocabulary) = MorseModel(
    WordEncoder(o,v), ContextEncoder(o,v), OutputEncoder(o,v),
    Embed(input=length(v.tags), output=o[:embedSizes][2]),
    Decoder(o,v),
    Linear(input=o[:hiddenSizes][3], output=length(v.tags)),
    CrossEntropyLoss(dims=1)
)

function loss(M::MorseModel, d::SentenceBatch; v::Vocabulary)
    # Word Encoder
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=true)
    # Context Encoder
    hiddens1 = M.contextEncoder(hiddens2[1]; training=true)
    # output encoder
    esize = size(M.outputEmbed.weight,1);
    tagEmbeddings = [reshape(M.outputEmbed(d.seqOutputs[i,range]),esize,1,length(range)) for (i,range)=enumerate(d.tagRange)]
    hiddens3 = M.outputEncoder(tagEmbeddings)
    # Decoder
    hiddens2 = hiddens2 .+ hiddens3
    recurrentLoss(M, d, hiddens1, hiddens2; startIndex=v.specialIndices.bow)/sum(d.masks)
end

function predict(M::MorseModel, d::SentenceBatch; v::Vocabulary)
    # Word Encoder
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=false)
    # Context Encoder
    hiddens1 = M.contextEncoder(hiddens2[1]; training=false)
    # Decoder
    total = 0.0
    wordNumber, timeSteps = size(d.seqOutputs)
    preds = zeros(Int, wordNumber, timeSteps)
    hiddens3  = (zero(M.outputEncoder.hidden), zero(M.outputEncoder.hidden))

    outEmbeddings = [];
    for w=1:wordNumber
        dh2 = (hiddens2[1][:,w:w] .+ hiddens3[1], hiddens2[2][:,w:w] .+ hiddens3[2])
        dh1 = (hiddens1[1][:,w:w], hiddens1[2][:,w:w])

        input = v.specialIndices.bow
        outEmbedding = [];
        continueStoring = true

        for t=1:timeSteps
            rnninput = M.outputEmbed([input])
            if continueStoring && input == v.specialIndices.eow
                continueStoring = false
            elseif continueStoring && input ∉ v.specialIndices && length(v.tags[input]) > 1
                push!(outEmbedding,rnninput)
            end
            dh1, dh2 = M.decoder(rnninput, dh1, dh2; training=false)
            h2 = dh2[1]
            scores = M.output(reshape(h2, size(h2,1), size(h2,2)))
            #prediction
            input = argmax(convert(Array,scores))[1]
            preds[w,t] = input
            # loss calculation
            total += M.loss(scores, d.seqOutputs[w:w,t] .* d.masks[w:w,t]; average=false)
        end

        if length(outEmbedding) != 0

            embedding = hcat(outEmbedding...)
            push!(outEmbeddings, reshape(embedding, size(embedding,1), 1, size(embedding,2)))

            if length(outEmbeddings) > M.outputEncoder.previous
                popfirst!(outEmbeddings);
            end

            if length(outEmbeddings) > 0
                hiddens3  = (zero(M.outputEncoder.hidden), zero(M.outputEncoder.hidden))
                for embs in outEmbeddings
                    hiddens3 = M.outputEncoder(embs, hiddens3[1], hiddens3[2])
                end
            end
        end
    end
    return preds, total
end

(M::MorseModel)(sentence::String; v::Vocabulary,p::Parser, maxL=20) = predict(M,sentence;v=v,p=p,maxL=maxL)

"""
    S2S <: Sequential
    S2S(o::Dict,v::Vocabulary)
    loss(M::S2S, d::SentenceBatch; v::Vocabulary)
    predict(M::S2S, d::SentenceBatch; v::Vocabulary)

    Ablation of `MorseModel` where `ContextEncoder` and `OutputEncoder` removed.
"""
struct S2S <: Sequential
    wordEncoder::WordEncoder
    outputEmbed::Embed
    decoder::Decoder
    output::Linear
    loss::CrossEntropyLoss
end

S2S(o::Dict,v::Vocabulary) = S2S(
    WordEncoder(o,v),
    Embed(input=length(v.tags), output=o[:embedSizes][2]),
    Decoder(o,v),
    Linear(input=o[:hiddenSizes][3], output=length(v.tags)),
    CrossEntropyLoss(dims=1)
)

function loss(M::S2S, d::SentenceBatch; v::Vocabulary)
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=true)
    recurrentLoss(M, d, (nothing, nothing), hiddens2; startIndex=v.specialIndices.bow) / sum(d.masks)
end

function predict(M::S2S, d::SentenceBatch; v::Vocabulary)
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=false)
    recurrentPredict(M, d, (nothing, nothing), hiddens2; startIndex=v.specialIndices.bow)
end

(M::S2S)(sentence::String; v::Vocabulary,p::Parser, maxL=20) = predict(M,sentence;v=v,p=p,maxL=maxL)
"""
    S2SContext <: Sequential
    S2SContext(config::Dict,v::Vocabulary)
    loss(M::S2SContext, d::SentenceBatch; v::Vocabulary)
    predict(M::S2SContext, d::SentenceBatch; v::Vocabulary)

    Ablation of the `MorseModel` where `OutputEncoder` removed
"""
struct S2SContext <: Sequential
    wordEncoder::WordEncoder
    contextEncoder::ContextEncoder
    outputEmbed::Embed
    decoder::Decoder
    output::Linear
    loss::CrossEntropyLoss
end

S2SContext(o::Dict,v::Vocabulary) = S2SContext(
    WordEncoder(o,v),
    ContextEncoder(o,v),
    Embed(input=length(v.tags), output=o[:embedSizes][2]),
    Decoder(o,v),
    Linear(input=o[:hiddenSizes][3], output=length(v.tags)),
    CrossEntropyLoss(dims=1)
)

function loss(M::S2SContext, d::SentenceBatch; v::Vocabulary)
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=true)
    hiddens1 = M.contextEncoder(hiddens2[1]; training=true)
    recurrentLoss(M, d, hiddens1, hiddens2; startIndex=v.specialIndices.bow) / sum(d.masks)
end

function predict(M::S2SContext, d::SentenceBatch; v::Vocabulary)
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=false)
    hiddens1 = M.contextEncoder(hiddens2[1]; training=false)
    recurrentPredict(M, d, hiddens1, hiddens2; startIndex=v.specialIndices.bow)
end

(M::S2SContext)(sentence::String; v::Vocabulary,p::Parser, maxL=20) = predict(M,sentence;v=v,p=p,maxL=maxL)
#####
##### Discriminative Models
#####

"""
    Classifier <: Discriminative
    Classifier(config::Dict,v::Vocabulary)
    loss(M::Classifier, d::SentenceBatch; v::Vocabulary)
    predict(M::Classifier, d::SentenceBatch; v::Vocabulary)

    Ablation of MorseModel where `OutputEncoder` removed and `Decoder` replaced with a classifer.
"""
struct Classifier <: Discriminative
    wordEncoder::WordEncoder
    contextEncoder::ContextEncoder
    output::Linear
    loss::CrossEntropyLoss
end

Classifier(o::Dict,v::Vocabulary) = Classifier(
    WordEncoder(o,v),
    ContextEncoder(o,v),
    Linear(input=o[:hiddenSizes][3]+o[:hiddenSizes][1], output=length(v.comptags)),
    CrossEntropyLoss(dims=1)
)

function loss(M::Classifier, d::SentenceBatch; v::Vocabulary)
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=true)
    hiddens1 = M.contextEncoder(hiddens2[1]; training=true)
    scores = M.output(vcat(hiddens1[1],hiddens2[1]))
    M.loss(scores,d.compositeOutputs)
end

function predict(M::Classifier, d::SentenceBatch; v::Vocabulary)
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=false)
    hiddens1 = M.contextEncoder(hiddens2[1]; training=false)
    scores   = M.output(vcat(hiddens1[1],hiddens2[1]))
    preds    = vec(getindex.(argmax(convert(Array,scores),dims=1),1))
    loss     = M.loss(scores,d.compositeOutputs)
    return preds,loss
end

(M::Classifier)(sentence::String; v::Vocabulary,p::Parser, maxL=20) = predict(M,sentence;v=v,p=p,maxL=maxL)
#####
##### Disambiguator Models
#####

abstract type Disambiguator <: Sequential; end;

struct S2SContexDis <: Disambiguator
    wordEncoder::WordEncoder
    contextEncoder::ContextEncoder
    outputEmbed::Embed
    decoder::Decoder
    output::Linear
    loss::CrossEntropyLoss
end

S2SContexDis(o::Dict,v::Vocabulary) = S2SContexDis(
    WordEncoder(o,v),
    ContextEncoder(o,v),
    Embed(input=length(v.tags), output=o[:embedSizes][2]),
    Decoder(o,v),
    Linear(input=o[:hiddenSizes][3], output=length(v.tags)),
    CrossEntropyLoss(dims=1)
)

function loss(M::S2SContexDis, d::SentenceBatch; v::Vocabulary)
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=true)
    hiddens1 = M.contextEncoder(hiddens2[1]; training=true)
    recurrentLoss(M, d, hiddens1, hiddens2; startIndex=v.specialIndices.bow) / sum(d.masks)
end

function predict(M::S2SContexDis, d::SentenceBatch; v::Vocabulary)
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=false)
    hiddens1 = M.contextEncoder(hiddens2[1]; training=true)
    recurrentPredict(M, d, hiddens1, hiddens2; v=v)
end

"""
    recurrentPredict(M::Disambiguator, d::SentenceBatch, hiddens1, hiddens2; startIndex=4)

    Make prediciton with recurrent decoder for disambiguation models except MorseDis
"""
function recurrentPredict(M::Disambiguator, d::SentenceBatch, hiddens1, hiddens2; v::Vocabulary)

    preds = Vector{Int}[]
    total = 0.0

    for (w,encodedIO) in enumerate(d.encodedIOs)

        min_loss = typemax(Float64)
        min_loss_seq = Int[];

        for (i,analysis) in enumerate(encodedIO.analyses)

            h1 = (hiddens1[1][:,w:w],hiddens1[2][:,w:w])
            h2 = (hiddens2[1][:,w:w],hiddens2[2][:,w:w])

            seq = [analysis.lemma; analysis.tags; v.specialIndices.eow]

            wtotal = 0.0
            input = v.specialIndices.bow

            for o in seq
                h1, h2 = M.decoder(M.outputEmbed([input]), h1, h2; training=false)
                hh2 = h2[1]
                scores = M.output(reshape(hh2,size(hh2,1),size(hh2,2)))
                wtotal += M.loss(scores,[o])
                input   = o
            end

            loss = wtotal
            if loss < min_loss
                min_loss = loss; min_loss_seq=seq;
            end
        end
        push!(preds,min_loss_seq)
        total+=min_loss
    end

    return PadSequenceArray(preds; pad=v.specialIndices.mask),total
    #TO-DO: total causes incorrect loos in evaluate.jl
end

"""
    MorseDis <: Disambiguator
    MorseDis(config::Dict,v::Vocabulary)
    loss(M::MorseDis, d::SentenceBatch; v::Vocabulary)
    predict(M::MorseDis, d::SentenceBatch; v::Vocabulary)
"""
struct MorseDis <: Disambiguator
    wordEncoder::WordEncoder
    contextEncoder::ContextEncoder
    outputEncoder::OutputEncoder
    outputEmbed::Embed
    decoder::Decoder
    output::Linear
    loss::CrossEntropyLoss
end

MorseDis(o::Dict,v::Vocabulary) = MorseDis(
    WordEncoder(o,v), ContextEncoder(o,v), OutputEncoder(o,v),
    Embed(input=length(v.tags), output=o[:embedSizes][2]),
    Decoder(o,v),
    Linear(input=o[:hiddenSizes][3], output=length(v.tags)),
    CrossEntropyLoss(dims=1)
)

function loss(M::MorseDis, d::SentenceBatch; v::Vocabulary)
    # Word Encoder
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=true)
    # Context Encoder
    hiddens1 = M.contextEncoder(hiddens2[1]; training=true)
    # output encoder
    esize = size(M.outputEmbed.weight,1);
    tagEmbeddings = [reshape(M.outputEmbed(d.seqOutputs[i,range]),esize,1,length(range)) for (i,range)=enumerate(d.tagRange)]
    hiddens3 = M.outputEncoder(tagEmbeddings)
    # Decoder
    hiddens2 = hiddens2 .+ hiddens3
    recurrentLoss(M, d, hiddens1, hiddens2; startIndex=v.specialIndices.bow)/sum(d.masks)
end

function  predict(M::MorseDis, d::SentenceBatch; v::Vocabulary)
    # Word Encoder
    hiddens2 = M.wordEncoder(d.seqInputs, d.seqSizes, d.unsortingIndices; training=false)
    # Context Encoder
    hiddens1 = M.contextEncoder(hiddens2[1]; training=false)
    # Output Encoder
    hiddens3  = (zero(M.outputEncoder.hidden), zero(M.outputEncoder.hidden))
    # Decoder

    outEmbeddings = [];
    preds = Vector{Int}[]
    total = 0.0

    for (w,encodedIO) in enumerate(d.encodedIOs)

        min_loss = typemax(Float64)
        min_loss_seq::Union{Nothing,Vector{Int}} = nothing;
        min_outEmbed::Union{Nothing,Vector} = nothing;

        for (i,analysis) in enumerate(encodedIO.analyses)

            h2 = (hiddens2[1][:,w:w] .+ hiddens3[1], hiddens2[2][:,w:w] .+ hiddens3[2])
            h1 = (hiddens1[1][:,w:w],hiddens1[2][:,w:w])

            wtotal = 0.0
            input = v.specialIndices.bow
            outEmbedding = [];
            seq = [analysis.lemma; analysis.tags; v.specialIndices.eow]

            for o in seq
                rnninput = M.outputEmbed([input])
                 if input ∉ v.specialIndices && length(v.tags[input])>1
                    push!(outEmbedding,rnninput)
                end
                h1, h2 = M.decoder(rnninput, h1, h2; training=false)
                hh2 = h2[1]
                scores = M.output(reshape(hh2,size(hh2,1),size(hh2,2)))
                wtotal += M.loss(scores,[o])
                input   = o
            end

            loss = wtotal
            if loss < min_loss
                min_loss = loss; min_loss_seq=seq; min_outEmbed=outEmbedding;
            end
        end

        push!(preds,min_loss_seq)
        total += min_loss

        if length(min_outEmbed) != 0

            embedding = hcat(min_outEmbed...)
            push!(outEmbeddings, reshape(embedding, size(embedding,1), 1, size(embedding,2)))

            if length(outEmbeddings) > M.outputEncoder.previous
                popfirst!(outEmbeddings);
            end

            if length(outEmbeddings) > 0
                hiddens3  = (zero(M.outputEncoder.hidden), zero(M.outputEncoder.hidden))
                for embs in outEmbeddings
                    hiddens3 = M.outputEncoder(embs, hiddens3[1], hiddens3[2])
                end
            end
        end
    end
    return PadSequenceArray(preds; pad=v.specialIndices.mask),total
end

#####
##### Utils
#####
function fgbias!(m::LSTM)
    for inputSource in (:i,:h), layer in (:10,:11)
        bias = get(m.gatesview, Symbol(:b,inputSource,:f,layer), nothing)
        if bias !== nothing
            bias .= 0.5
        end
    end
    return m
end

function transfer!(target, tv::Vocabulary, source, sv::Vocabulary)
    for f in (:wordEncoder, :decoder, :contextEncoder, :outputEmbed, :output, :outputEncoder)
        if isdefined(target,f) && isdefined(source,f)
            transfer!(getproperty(target,f), tv, getproperty(source,f), sv)
        end
    end
end

function transfer!(target::Embed, tv::Vocabulary, source::Embed, sv::Vocabulary)
    transfer!(target.weight, tv.tags, source.weight, sv.tags; axis=2)
end

function transfer!(target::Linear, tv::Vocabulary, source::Linear, sv::Vocabulary)
    transfer!(target.mult.weight, tv.tags, source.mult.weight, sv.tags; axis=1)
    transfer!(target.bias, tv.tags, source.bias, sv.tags; axis=1)
end

function transfer!(target, tv::IndexedDict, source, sv::IndexedDict; axis=1)
    for (k,i) in sv.toIndex
        if (j=get(tv,k,nothing)) !== nothing
            if ndims(target) == 2
                if axis == 1
                    target[j,:] = source[i,:]
                else
                    target[:,j] = source[:,i]
                end
            else
                 target[j] = source[i]
            end
        end
    end
end

function transfer!(source::Decoder, sv::Vocabulary, target::Decoder, tv::Vocabulary)
    transfer!(source.L1,target.L1)
    transfer!(source.L2,target.L2)
end

function transfer!(source::ContextEncoder,sv::Vocabulary,target::ContextEncoder,tv::Vocabulary)
    transfer!(source.encoder,target.encoder)
    transfer!(source.reducer,target.reducer)
end

function transfer!(source::OutputEncoder,sv::Vocabulary,target::OutputEncoder,tv::Vocabulary)
    transfer!(source.encoder,target.encoder)
end

function transfer!(target::WordEncoder, tv::Vocabulary, source::WordEncoder, sv::Vocabulary)
    transfer!(target.embed.weight, tv.chars, source.embed.weight, sv.chars; axis=2)
    transfer!(target.encoder, source.encoder)
end

function transfer!(target::Linear, source::Linear)
    transfer!(target.mult,source.mult)
    transfer!(target.bias,source.bias)
end
transfer!(target::Multiply, source::Multiply) = transfer!(target.weight,source.weight)
transfer!(target::Dense, source::Dense) = transfer!(target.linear,source.linear)
transfer!(target::LSTM, source::LSTM) = transfer!(target.params,source.params)
transfer!(target::Param, source::Param) = copyto!(value(target),value(source))


using Test
using JPLEphemeris

using AstroBase

const AU = astronomical_unit(Float64, KiloMeter())[1]

const SPK_URL = Dict{Int, String}(
    430 => "https://naif.jpl.nasa.gov/pub/naif/generic_kernels/spk/planets/de430.bsp",
    405 => "https://naif.jpl.nasa.gov/pub/naif/generic_kernels/spk/planets/a_old_versions/de405.bsp",
)
const TEST_URL = Dict{Int, String}(
    430 => "ftp://ssd.jpl.nasa.gov/pub/eph/planets/test-data/430/testpo.430",
    405 => "ftp://ssd.jpl.nasa.gov/pub/eph/planets/test-data/testpo.405",
)

path = abspath(joinpath(dirname(@__FILE__), "..", "deps"))

for denum in (430, 405)
    if !isdir(path)
        mkdir(path)
    end
    if !isfile("$path/de$denum.bsp")
        download(SPK_URL[denum], "$path/de$denum.bsp")
    end
    if !isfile("$path/testpo.$denum")
        download(TEST_URL[denum], "$path/testpo.$denum")
    end
end

@testset "JPLEphemeris" begin
    include("basic.jl")
    include("kernels.jl")
end


#
# Powell.jl --
#
# Mike Powell's derivative free optimization methods for Julia.
#
# -----------------------------------------------------------------------------
#
# This file is part of OptimPack.jl which is licensed under the MIT
# "Expat" License:
#
# Copyright (C) 2015-2019, Éric Thiébaut <https://github.com/emmt/OptimPack.jl>.
#

module Powell

export
    iterate,
    restart,
    getstatus,
    getreason,
    getradius,
    getncalls,
    getlastf,
    Cobyla,
    cobyla,
    cobyla!,
    Newuoa,
    newuoa,
    newuoa!,
    Bobyqa,
    bobyqa,
    bobyqa!

using Compat

import Base: ==, iterate

# The dynamic libraries implementing the methods.
import ..libcobyla
import ..libbobyqa
import ..libnewuoa

abstract type AbstractContext end
abstract type AbstractStatus end

==(a::T, b::T) where {T<:AbstractStatus} = (a._code == b._code)
==(a::AbstractStatus, b::AbstractStatus) = false

"""
The `iterate(ctx, ...)` method performs the next iteration of the reverse
communication associated with the context `ctx`.  Other arguments depend on the
type of algorithm.

For **COBYLA** algorithm, the next iteration is performed by:

    iterate(ctx, f, x, c) -> status

or

    iterate(ctx, f, x) -> status

on entry, the workspace status must be `COBYLA_ITERATE`, `f` and `c` are the
function value and the constraints at `x`, the latter can be omitted if there
are no constraints.  On exit, the returned value (the new workspace status) is:
`COBYLA_ITERATE` if a new trial point has been stored in `x` and if user is
requested to compute the function value and the constraints on the new point;
`COBYLA_SUCCESS` if algorithm has converged and `x` has been set with the
variables at the solution (the corresponding function value can be retrieved
with `getlastf`); anything else indicates an error (see `getreason`
for an explanatory message).

For **NEWUOA** algorithm, the next iteration is performed by:

    iterate(ctx, f, x) -> status

on entry, the wokspace status must be `NEWUOA_ITERATE`, `f` is the function
value at `x`.  On exit, the returned value (the new wokspace status) is:
`NEWUOA_ITERATE` if a new trial point has been stored in `x` and if user is
requested to compute the function value for the new point; `NEWUOA_SUCCESS` if
algorithm has converged; anything else indicates an error (see `getreason` for
an explanatory message).

"""
function iterate end

"""

    restart(ctx) -> status

restarts the reverse communication algorithm associated with the context `ctx`
using the same parameters.  The return value is the new status of the
algorithm, see `getstatus` for details.

"""
function restart end

"""

    getstatus(ctx) -> status

get the current status of the reverse communication algorithm associated with
the context `ctx`.  Possible values are:

* for **COBYLA**: `COBYLA_ITERATE`, if user is requested to compute `f(x)` and
  `c(x)`; `COBYLA_SUCCESS`, if algorithm has converged;

* for **NEWUOA**: `NEWUOA_ITERATE`, if user is requested to compute `f(x)`;
  `NEWUOA_SUCCESS`, if algorithm has converged;

Anything else indicates an error (see `getreason` for an explanatory message).

"""
function getstatus end

"""

    getreason(ctx) -> msg

or

    getreason(status) -> msg

get an explanatory message about the current status of the reverse
communication algorithm associated with the context `ctx` or with the status
returned by an optimization method of by `getstatus(ctx)`.

"""
function getreason end

getreason(ctx::AbstractContext) = getreason(getstatus(ctx))

"""

    getlastf(ctx) -> fx

get the last function value in the reverse communication algorithm associated
with the context `ctx`.  Upon convergence of `iterate`, this value corresponds
to the function at the solution; otherwise, this value corresponds to the
previous set of variables.

"""
function getlastf end

"""

    getncalls(ctx) -> nevals

get the current number of function evaluations in the reverse communication
algorithm associated with the context `ctx`.  Result is -1 if something is
wrong, nonnegative otherwise.

"""
function getncalls end

"""

    getradius(ctx) -> rho

get the current size of the trust region of the reverse communication algorithm
associated with the context `ctx`.  Result is 0 if algorithm not yet started
(before first iteration), -1 if something is wrong, strictly positive
otherwise.

"""
function getradius end

"""

```julia
grow!(x, n) -> x
```

grows vector `x` so that it has at least `n` elements, does nothing if `x` is
large enough.  Argument `x` is returned.

See also [`resize!`](@ref).

"""
grow!(x::Vector, n::Integer) = begin
    length(x) < n && resize!(x, n)
    return x
end

include("newuoa.jl")
import .Newuoa: newuoa, newuoa!

include("cobyla.jl")
import .Cobyla: cobyla, cobyla!

include("bobyqa.jl")
import .Bobyqa: bobyqa, bobyqa!

end # module Powell


""" Data structure for undirected wiring diagrams.
"""
module UndirectedWiringDiagrams
export UndirectedWiringDiagram, UntypedUWD, TypedUWD,
  outer_box, box, junction, nboxes, njunctions, boxes, junctions,
  ports, ports_with_junction, junction_type, port_type,
  add_box!, add_junction!, add_junctions!, set_junction!, add_wire!,
  add_wires!, singleton_diagram, cospan_diagram, junction_diagram,
  ocompose, substitute

using ...Present, ...CategoricalAlgebra.CSets, ...CategoricalAlgebra.Limits
using ...CategoricalAlgebra.FinSets: FinSet, FinFunction
using ...Theories: dom, codom, compose, ⋅, id, oplus
import ..DirectedWiringDiagrams: box, boxes, nboxes, add_box!, add_wire!,
  add_wires!, singleton_diagram, ocompose, substitute

# Data types
############

@present TheoryUWD(FreeSchema) begin
  Box::Ob
  Port::Ob
  OuterPort::Ob
  Junction::Ob

  box::Hom(Port, Box)
  junction::Hom(Port, Junction)
  outer_junction::Hom(OuterPort, Junction)
end

const UndirectedWiringDiagram = const AbstractUWD = AbstractACSetType(TheoryUWD)
const UntypedUWD = CSetType(TheoryUWD, index=[:box, :junction, :outer_junction])

@present TheoryTypedUWD <: TheoryUWD begin
  Type::Data

  port_type::Attr(Port, Type)
  outer_port_type::Attr(OuterPort, Type)
  junction_type::Attr(Junction, Type)

  compose(junction, junction_type) == port_type
  compose(outer_junction, junction_type) == outer_port_type
end

const AbstractTypedUWD = AbstractACSetType(TheoryTypedUWD)
const TypedUWD = ACSetType(TheoryTypedUWD,
                           index=[:box, :junction, :outer_junction])

function (::Type{UWD})(nports::Int) where UWD <: AbstractUWD
  d = UWD()
  add_parts!(d, :OuterPort, nports)
  return d
end

function (::Type{UWD})(port_types::AbstractVector) where UWD <: AbstractUWD
  d = UWD()
  nports = length(port_types)
  add_parts!(d, :OuterPort, nports, outer_port_type=port_types)
  return d
end

UndirectedWiringDiagram(nports::Int) = UntypedUWD(nports)
UndirectedWiringDiagram(port_types::AbstractVector{T}) where T =
  TypedUWD{T}(port_types)

# Imperative interface
######################

# Accessors
#----------

outer_box(::AbstractACSet) = 0
box(d::AbstractACSet, args...) = subpart(d, args..., :box)
junction(d::AbstractACSet, args...; outer::Bool=false) =
  subpart(d, args..., outer ? :outer_junction : :junction)

function junction(d::AbstractACSet, port::Tuple{Int,Int})
  box, nport = port
  box == outer_box(d) ?
    junction(d, nport, outer=true) : junction(d, ports(d, box)[nport])
end

nboxes(d::AbstractACSet) = nparts(d, :Box)
njunctions(d::AbstractACSet) = nparts(d, :Junction)
boxes(d::AbstractACSet) = parts(d, :Box)
junctions(d::AbstractACSet) = parts(d, :Junction)

ports(d::AbstractACSet; outer::Bool=false) =
  parts(d, outer ? :OuterPort : :Port)
ports(d::AbstractACSet, box) =
  box == outer_box(d) ? parts(d, :OuterPort) : incident(d, box, :box)
ports_with_junction(d::AbstractACSet, junction; outer::Bool=false) =
  incident(d, junction, outer ? :outer_junction : :junction)

junction_type(d::AbstractACSet, args...) = subpart(d, args..., :junction_type)
port_type(d::AbstractACSet, args...; outer::Bool=false) =
  subpart(d, args..., outer ? :outer_port_type : :port_type)

function port_type(d::AbstractACSet, port::Tuple{Int,Int})
  box, nport = port
  box == outer_box(d) ?
    port_type(d, nport, outer=true) : port_type(d, ports(d, box)[nport])
end

# Mutators
#---------

add_box!(d::AbstractUWD; data...) = add_part!(d, :Box; data...)

function add_box!(d::AbstractUWD, nports::Int; data...)
  box = add_box!(d; data...)
  ports = add_parts!(d, :Port, nports, box=fill(box, nports))
  box
end

function add_box!(d::AbstractUWD, port_types::AbstractVector; data...)
  box = add_box!(d; data...)
  nports = length(port_types)
  ports = add_parts!(d, :Port, nports, box=fill(box, nports),
                     port_type=port_types)
  box
end

add_junction!(d::AbstractUWD; data...) = add_part!(d, :Junction; data...)
add_junction!(d::AbstractUWD, type; data...) =
  add_part!(d, :Junction; junction_type=type, data...)

add_junctions!(d::AbstractUWD, njunctions::Int; data...) =
  add_parts!(d, :Junction, njunctions; data...)
add_junctions!(d::AbstractUWD, types::AbstractVector; data...) =
  add_parts!(d, :Junction, length(types); junction_type=types, data...)

function set_junction!(d::AbstractUWD, port, junction; outer::Bool=false)
  if has_subpart(d, :junction_type)
    ptype, jtype = port_type(d, port, outer=outer), junction_type(d, junction)
    all(ptype .== jtype) || error(
      "Domain error: port type $ptype and junction type $jtype do not match")
  end
  set_subpart!(d, port, outer ? :outer_junction : :junction, junction)
end
set_junction!(d::AbstractUWD, junction; kw...) =
  set_junction!(d, :, junction; kw...)

function set_junction!(d::AbstractUWD, port::Tuple{Int,Int}, junction)
  box, nport = port
  if box == outer_box(d)
    set_junction!(d, nport, junction, outer=true)
  else
    set_junction!(d, ports(d, box)[nport], junction)
  end
end

""" Wire together two ports in an undirected wiring diagram.

A convenience method that creates and sets junctions as needed. Ports are only
allowed to have one junction, so if both ports already have junctions, then the
second port is assigned the junction of the first. The handling of the two
arguments is otherwise symmetric.

FIXME: When both ports already have junctions, the two junctions should be
*merged*. To do this, we must implement `merge_junctions!` and thus also
`rem_part!`.
"""
function add_wire!(d::AbstractUWD, port1::Tuple{Int,Int}, port2::Tuple{Int,Int})
  j1, j2 = junction(d, port1), junction(d, port2)
  if j1 > 0
    set_junction!(d, port2, j1)
  elseif j2 > 0
    set_junction!(d, port1, j2)
  else
    j = has_subpart(d, :junction_type) ?
      add_junction!(d, port_type(d, port1)) : add_junction!(d)
    set_junction!(d, port1, j)
    set_junction!(d, port2, j)
  end
end
add_wire!(d, wire::Pair) = add_wire!(d, first(wire), last(wire))

function add_wires!(d::AbstractUWD, wires)
  for wire in wires
    add_wire!(d, wire)
  end
end

# Other constructors
#-------------------

function singleton_diagram(::Type{UWD}, port_types;
                           data...) where UWD <: AbstractUWD
  d = UWD(port_types)
  add_box!(d, port_types; data...)
  js = add_junctions!(d, port_types)
  set_junction!(d, js)
  set_junction!(d, js, outer=true)
  return d
end

""" Undirected wiring diagram defined by a cospan.

The wiring diagram has a single box. The ports of this box, the outer ports, the
junctions, and the connections between them are defined by the cospan. Thus,
this function generalizes [`singleton_diagram`](@ref).

See also: [`junction_diagram`](@ref).
"""
function cospan_diagram(::Type{UWD}, inner::FinFunction{Int,Int},
                        outer::FinFunction{Int,Int}, junction_types=nothing;
                        data...) where UWD <: AbstractUWD
  @assert codom(inner) == codom(outer)
  if isnothing(junction_types)
    junction_types = length(codom(inner))
  end
  cospan_diagram(UWD, collect(inner), collect(outer), junction_types; data...)
end
function cospan_diagram(::Type{UWD}, inner::AbstractVector{Int},
                        outer::AbstractVector{Int}, junction_types;
                        data...) where UWD <: AbstractUWD
  d = UWD(map_port_types(junction_types, outer))
  add_box!(d, map_port_types(junction_types, inner); data...)
  add_junctions!(d, junction_types)
  set_junction!(d, inner)
  set_junction!(d, outer, outer=true)
  return d
end

""" Undirected wiring diagram with no boxes, only junctions.

See also: [`singleton_diagram`](@ref), [`cospan_diagram`](@ref).
"""
function junction_diagram(::Type{UWD}, outer::FinFunction{Int,Int},
                          junction_types=nothing) where UWD <: AbstractUWD
  if isnothing(junction_types)
    junction_types = length(codom(outer))
  end
  junction_diagram(UWD, collect(outer), junction_types)
end
function junction_diagram(::Type{UWD}, outer::AbstractVector{Int},
                          junction_types) where UWD <: AbstractUWD
  d = UWD(map_port_types(junction_types, outer))
  add_junctions!(d, junction_types)
  set_junction!(d, outer, outer=true)
  return d
end

map_port_types(::Int, g) = length(g)
map_port_types(types::AbstractVector, g) = types[g]

# Operadic interface
####################

function ocompose(f::AbstractUWD, gs::AbstractVector{<:AbstractUWD})
  substitute(f, boxes(f), gs)
end
function ocompose(f::AbstractUWD, i::Int, g::AbstractUWD)
  substitute(f, [i], [g])
end

function substitute(f::UWD, indices::AbstractVector{Int},
                    gs::AbstractVector{UWD}) where {UWD <: AbstractUWD}
  @assert length(indices) == length(gs)
  h = UWD()
  copy_parts!(h, f, OuterPort=:)

  # Copy boxes from original diagram and diagrams to be substituted.
  gs_index = Vector{Union{Int,Nothing}}(nothing, nboxes(f))
  gs_index[indices] = 1:length(indices)
  for (i,j) in enumerate(gs_index)
    if isnothing(j)
      copy_parts!(h, f, Box=[i], Port=ports(f, i))
    else
      g = gs[j]
      copy_parts!(h, g, Box=boxes(g), Port=flat(ports(g, boxes(g))))
    end
  end

  # Add junctions obtained from pushout of attributed junction sets.
  Junctions = ACSetTableType(UWD, :Junction)
  f_junctions, g_junctions, apex = Junctions(), Junctions(), Junctions()
  copy_parts_only!(f_junctions, f)
  for g in gs
    copy_parts_only!(g_junctions, g)
    copy_parts_only!(apex, g, Junction=junction(g, outer=true))
  end
  f_inner = ACSetTransformation(apex, f_junctions, Junction=flat(
    junction(f, ports(f, i)) for (i,j) in enumerate(gs_index) if !isnothing(j)))
  g_outer = ACSetTransformation(apex, g_junctions, Junction=oplus([
    FinFunction(junction(g, outer=true), njunctions(g)) for g in gs ]))
  colim = pushout(f_inner, g_outer)
  copy_parts_only!(h, ob(colim))

  # Assign junctions to all ports.
  f_leg, g_leg = (leg[:Junction] for leg in colim)
  gs_incl = legs(coproduct([ FinSet(njunctions(g)) for g in gs ]))
  set_junction!(h, map(f_leg, junction(f, outer=true)), outer=true)
  set_junction!(h, flatmap(enumerate(gs_index)) do (i,j)
    if isnothing(j)
      map(f_leg, junction(f, ports(f, i)))
    else
      map(gs_incl[j] ⋅ g_leg, junction(gs[j]))
    end
  end)
  return h
end

flat(x) = reduce(vcat, x, init=Int[])
flatmap(f, xs...) = mapreduce(f, vcat, xs..., init=Int[])

end


module UnitSystems

#   This file is part of UnitSystems.jl. It is licensed under the MIT license
#   UnitSystems Copyright (C) 2020 Michael Reed

import Base: @pure, length, time

const Systems = (:Metric,:SI2019,:CODATA,:Conventional,:MTS,:English,:EnglishUS,:IAU,:SI1976,:Mixed,:ESU2019,:EMU2019,:EMU,:ESU,:Gauss,:LorentzHeaviside,:Thomson,:Kennelly,:Planck,:PlanckGauss,:Stoney,:Hartree,:Rydberg,:Schrodinger,:Electronic,:Natural,:NaturalGauss,:QCD,:QCDGauss,:QCDoriginal)
const Constants = (:hyperfine,:lightspeed,:planck,:planckreduced,:electronmass,:molarmass,:boltzmann,:permeability,:rationalization,:lorentz,:luminousefficacy)
const Physics = (:atomicmass,:protonmass,:planckmass,:newton,:einstein,:hartree,:rydberg,:bohr,:bohrreduced,:electronradius,:avogadro,:universal,:stefan,:radiationdensity,:permittivity,:coulomb,:ampere,:biotsavart,:charge,:faraday,:impedance,:conductance,:klitzing,:josephson,:magneticflux,:magneton)
const Kinematic = (:time,:length,:area,:volume,:wavenumber,:fuelefficiency,:frequency,:frequencydrift,:speed,:acceleration,:jerk,:snap,:volumeflow)
const Mechanical = (:mass,:massflow,:lineardensity,:areadensity,:density,:specificvolume,:force,:stiffness,:pressure,:compressibility,:viscosity,:diffusivity,:rotationalinertia,:momentum,:angularmomentum,:yank,:energy,:specificenergy,:action,:fluence,:power,:powerdensity,:intensity,:spectralflux,:soundexposure,:impedance,:specificimpedance,:admittance,:compliance,:inertance)
const Electromagnetic = (:charge,:chargedensity,:linearchargedensity,:exposure,:mobility,:current,:currentdensity,:resistance,:conductance,:resistivity,:conductivity,:capacitance,:inductance,:reluctance,:permeance,:permittivity,:permeability,:susceptibility,:specificsusceptibility,:demagnetizingfactor,:vectorpotential,:electricpotential,:magneticpotential,:electricfield,:magneticfield,:electricflux,:magneticflux,:electricfluxdensity,:magneticfluxdensity,:electricdipolemoment,:magneticdipolemoment,:electricpolarizability,:magneticpolarizability,:magneticmoment,:magnetizability,:magnetization,:specificmagnetization,:rigidity,:polestrength)
const Thermodynamic = (:temperature,:entropy,:specificentropy,:volumeheatcapacity,:thermalconductivity,:thermalconductance,:thermalresistance,:thermalexpansion,:lapserate)
const Molar = (:molarmass,:molality,:mole,:molarity,:molarvolume,:molarentropy,:molarenergy,:molarconductivity,:molarsusceptibility,:catalysis,:specificity)
const Photometric = (:luminousflux,:luminance,:luminousenergy,:luminousexposure,:luminousefficacy)
const Mechanics = [Kinematic...,Mechanical...]
const Convert = [Mechanics...,Electromagnetic...,Thermodynamic...,Molar...,Photometric...]

listext(x) = join(x,"`, `")

@pure measure(x) = x

# universe

"""
    Coupling{αG,α,μₑᵤ,μₚᵤ}

Specification of `Universe` with the dimensionless `Coupling` constants `coupling`, `finestructure`, `electronunit`, `protonunit`, and `protonelectron`.
Alterations to these values can be facilitated and quantified using parametric polymorphism.
Due to the `Coupling` interoperability, the `MeasureSystems` package is made possible to support calculations with `Measurements` having error standard deviations.
"""
struct Coupling{αG,α,μₑᵤ,μₚᵤ} end
@pure coupling(U::Coupling{αG}) where αG = measure(αG)
@pure finestructure(U::Coupling{αG,α}) where {αG,α} = measure(α)
@pure electronunit(U::Coupling{αG,α,μₑᵤ}) where {αG,α,μₑᵤ} = measure(μₑᵤ)
@pure protonunit(U::Coupling{αG,α,μₑᵤ,μₚᵤ}) where {αG,α,μₑᵤ,μₚᵤ} = measure(μₚᵤ)
@pure protonelectron(U::Coupling) = protonunit(U)/electronunit(U)
Base.display(U::Coupling) = println("Coupling{αG=$(coupling(U)),α=$(finestructure(U)),μₑᵤ=$(electronunit(U)),μₚᵤ=$(protonunit(U))}")

# unit systems

"""
    UnitSystem{kB,ħ,𝘤,μ₀,mₑ,λ,αL}

Fundamental constants of physics are: `kB` Boltzmann's constant, `ħ` reduced Planck's constant, `𝘤` speed of light, `μ₀` vacuum permeability, `mₑ` electron rest mass, `λ` Gauss rationalization, and `αL` Lorentz's constant.
Primarily the `Metric` SI unit system is used in addition to the historic `English` engineering unit system.
These constants induce derived values for `avogadro`, `boltzmann`, `universal`, `planck`, `planckreduced`, `lightspeed`, `planckmass`, `atomicmass`, `protonmass`, `electronmass`, `newton`, `einstein`, `permeability`, `permittivity`, `coulomb`, and
additional constants `molarmass`, `hyperfine`, `luminousefficacy`, `stefan`, `radiationdensity`, `ampere`, `lorentz`, `biotsavart`, `rationalization`, `impedance`, `charge`, `magneton`, `conductance`, `faraday`, `magneticflux`, `josephson`, `klitzing`, `hartree`, `rydberg`, `bohr`, and `bohrreduced`.

Additional reference `UnitSystem` variants: `EMU`, `ESU`, `Gauss`, `LorentzHeaviside`, `MTS`, `SI2019`, `CODATA`, `Conventional`, `IAU`, `EnglishUS`; and natural atomic units based on gravitational coupling `αG` and the fine structure `1/αinv` constant (`Planck`, `PlanckGauss`, `Stoney`, `Hartree`, `Rydberg`, `Schrodinger`, `Electronic`, `Natural`, `NaturalGauss`, `QCD`, `QCDGauss`, and `QCDoriginal`).

**Derived unit conversions:**

Mechanics: `$(listext(Kinematic))`, `$(listext(Mechanical))`;
Electromagnetics: `$(listext(Electromagnetic))`;
Thermodynamics: `$(listext(Thermodynamic))`,
`$(listext(Molar))`, `$(listext(Photometric))`.
""" #`Rᵤ,mᵤ,σ,ħ,μ₀,ε₀,kₑ,𝘦,𝔉,RK,Z₀,G₀`
struct UnitSystem{kB,ħ,𝘤,μ₀,mₑ,λ,αL} end
@pure UnitSystem{k,ħ,𝘤,μ,m,λ}() where {k,ħ,𝘤,μ,m,λ} = UnitSystem{k,ħ,𝘤,μ,m,λ,1}()
@pure UnitSystem{k,ħ,𝘤,μ,m}() where {k,ħ,𝘤,μ,m} = UnitSystem{k,ħ,𝘤,μ,m,1}()
@pure boltzmann(::UnitSystem{k}) where k = measure(k)
@pure planckreduced(::UnitSystem{k,ħ}) where {k,ħ} = measure(ħ)
@pure lightspeed(::UnitSystem{k,ħ,𝘤}) where {k,ħ,𝘤} = measure(𝘤)
@pure permeability(::UnitSystem{k,ħ,𝘤,μ}) where {k,ħ,𝘤,μ} = measure(μ)
@pure electronmass(::UnitSystem{k,ħ,𝘤,μ,m}) where {k,ħ,𝘤,μ,m} = measure(m)
@pure rationalization(::UnitSystem{k,ħ,𝘤,μ,m,λ}) where {k,ħ,𝘤,μ,m,λ} = measure(λ)
@pure lorentz(::UnitSystem{k,ħ,𝘤,μ,m,λ,α}) where {k,ħ,𝘤,μ,m,λ,α} = measure(α)
# ΔνCs:s⁻¹, c:m⋅s⁻¹, h:kg⋅m²⋅s⁻¹, kB:kg⋅m²⋅s⁻²⋅K⁻¹, NA:mol⁻¹, Kcd: cd⋅sr⋅s³⋅kg⁻¹⋅m⁻²

isrationalized(U::UnitSystem) = rationalization(U) ≠ 4π

Base.display(U::UnitSystem) = println("UnitSystem{kB=$(boltzmann(U)),ħ=$(planckreduced(U)),𝘤=$(lightspeed(U)),μ₀=$(permeability(U)),mᵤ=$(electronmass(U)),λ=$(isrationalized(U) ? rationalization(U) : "4π"),αL=$(lorentz(U))}")

@pure universe(::UnitSystem) = Universe
@pure unit(x,y=1) = isapprox(y,x,rtol=eps()^0.9) ? y : x
@pure mass(U::UnitSystem,S::UnitSystem) = electronmass(U,S)
@pure electronmass(𝘩::Float64) = αinv^2*R∞*2𝘩/𝘤
@pure electronmass(𝘩::Float64,C::Coupling) = inv(finestructure(C))^2*R∞*2𝘩/𝘤
@pure planckmass(U::UnitSystem,C::Coupling=universe(U)) = electronmass(U,C)/√coupling(C)
@pure planck(U::UnitSystem,C::Coupling=universe(U)) = 2π*planckreduced(U,C)
@pure newton(U::UnitSystem,C::Coupling=universe(U)) = lightspeed(U,C)*planckreduced(U,C)/planckmass(U,C)^2
@pure charge(U::UnitSystem,C::Coupling=universe(U)) = sqrt(2planck(U)/(permeability(U)/finestructure(U))/(lightspeed(U)*rationalization(U)*lorentz(U)^2))

for unit ∈ (:coupling,:finestructure,:electronunit,:protonunit,:protonelectron)
    @eval @pure $unit(U::UnitSystem) = $unit(universe(U))
end
for unit ∈ (:boltzmann,:planckreduced,:lightspeed,:permeability,:electronmass,:molarmass)
    @eval @pure $unit(U::UnitSystem,C::Coupling) = $unit(U)
end
for unit ∈ Constants
    @eval @pure $unit(U::UnitSystem,S::UnitSystem) = unit($unit(S)/$unit(U))
end
for unit ∈ Convert
    @eval begin
        @pure @inline $unit(v::Real,U::UnitSystem) = $unit(v,U,Metric)
        @pure @inline $unit(v::Real,U::UnitSystem,S::UnitSystem) = (u=$unit(U,S);isone(u) ? v : v/u)
        @pure @inline $unit(v::Real,U::UnitSystem{kB,ħ,𝘤,μ₀,mₑ},S::UnitSystem{kB,ħ,𝘤,μ₀,mₑ}) where {kB,ħ,𝘤,μ₀,mₑ} = v
        @pure @inline $unit(U::UnitSystem{kB,ħ,𝘤,μ₀,mₑ},S::UnitSystem{kB,ħ,𝘤,μ₀,mₑ}) where {kB,ħ,𝘤,μ₀,mₑ} = 1
    end
    if unit ∉ (Constants...,:permittivity,:charge,:magneticflux,:impedance,:conductance)
        @eval @pure @inline $unit(U::UnitSystem) = $unit(U,Metric)
    end
end
for unit ∈ (Systems...,Constants...,Physics...,Convert...)
    @eval export $unit
end

function (U::UnitSystem)(JK=1,Js=1,ms=1,Hm=1,kg=1)
    kB = boltzmann(U)*JK
    ħ = planckreduced(U)*Js
    𝘤 = lightspeed(U)*ms
    μ₀ = permeability(U)*Hm
    mₑ = electronmass(U)*kg
    λ = rationalization(U)
    αL = lorentz(U)
    UnitSystem{kB,ħ,𝘤,μ₀,mₑ,λ,isone(αL) ? αL : αL/ms}()
end

# common conversion factors

const g₀,ft,ftUS = 9.80665,0.3048,1200/3937; const lbm,lbmUS = g₀/ft,g₀/ftUS
const slug,slugUS,rankine,kelvin,atm = 0.45359237lbm,0.45359237lbmUS,5/9,9/5,101325.0
const kcalₜₕ,kcal₄,kcal₁₀,kcal₂₀,kcalₘ,kcalᵢₜ = 4184,4204,4185.5,4182,4190,4186.8
const calₜₕ,cal₄,cal₁₀,cal₂₀,calₘ,calᵢₜ = (kcalₜₕ,kcal₄,kcal₁₀,kcal₂₀,kcalₘ,kcalᵢₜ)./1e3
const kcal = kcalₜₕ; const cal = kcal/1000 # calₜₕ thermal calorie

# fundamental constants, αinv = (34259-1/4366.8123)/250 # 137.036 exactly?

const ΔνCs,Kcd,mP = 9192631770.0,683*555.016/555,2.176434e-8 # planck mass (kg)
const NA,kB,𝘩,𝘤,𝘦 = 6.02214076e23,1.380649e-23,6.62607015e-34,299792458.,1.602176634e-19
const μₑᵤ,μₚᵤ,αinv,R∞ = 1/1822.888486209,1.007276466621,137.035999084,10973731.5681601
const mₑ,μ₀ = electronmass(𝘩),2𝘩/𝘤/αinv/𝘦^2 # ≈ 4π*(1e-7+5.5e-17), exact charge
const ħ,δμ₀,μₚₑ,Rᵤ,αL,αG = 𝘩/2π,μ₀-4π*1e-7,μₚᵤ/μₑᵤ,NA*kB,0.01/𝘤,(mₑ/mP)^2
const RK1990,RK2014,KJ1990,KJ2014 = 25812.807,25812.8074555,4.835979e14,4.835978525e14
const ħ1990,ħ2014 = 2/RK1990/KJ1990^2/π,2/RK2014/KJ2014^2/π
const mₑ1990,mₑ2014 = electronmass(2π*ħ1990),electronmass(2π*ħ2014)

# engineering units # Thomson: αL = 1/2

const Universe = Coupling{αG,1/αinv,μₑᵤ,μₚᵤ}()
const Gauss = UnitSystem{1e10*Rᵤ*mₑ/μₑᵤ,1e7*ħ,100𝘤,1,1000mₑ,4π,0.01/𝘤}()
const LorentzHeaviside = UnitSystem{1e10*Rᵤ*mₑ/μₑᵤ,1e7*ħ,100𝘤,1,1000mₑ,1,0.01/𝘤}()
const Thomson = UnitSystem{1e10*Rᵤ*mₑ/μₑᵤ,1e7*ħ,100𝘤,1,1000mₑ,4π,1/2}()
const Kennelly = UnitSystem{Rᵤ*mₑ/μₑᵤ/0.001,ħ,𝘤,1e-7,mₑ,4π}() # ?
const ESU = UnitSystem{1e10*Rᵤ*mₑ/μₑᵤ,1e7*ħ,100𝘤,(100𝘤)^-2,1000mₑ,4π}()
const ESU2019 = UnitSystem{1e7*kB,1e7*ħ,100𝘤,1e3*μ₀/𝘤^2,1000mₑ}()
const EMU = UnitSystem{1e10*Rᵤ*mₑ/μₑᵤ,1e7*ħ,100𝘤,1,1000mₑ,4π}()
const EMU2019 = UnitSystem{1e7*kB,1e7*ħ,100𝘤,1e7*μ₀,1000mₑ}()
const MTS = UnitSystem{1e6*Rᵤ*mₑ/μₑᵤ,1000ħ,𝘤,4π/1e4,mₑ/1000}()
const Mixed = UnitSystem{Rᵤ*mₑ/μₑᵤ/0.001,ħ,𝘤,μ₀,mₑ}()
const Metric = UnitSystem{Rᵤ*mₑ/μₑᵤ/0.001,ħ,𝘤,4π*1e-7,mₑ}()
const SI1976 = UnitSystem{8.31432mₑ/μₑᵤ/0.001,ħ,𝘤,4π*1e-7,mₑ}()
const SI2019 = UnitSystem{kB,ħ,𝘤,μ₀,mₑ}()
const CODATA = UnitSystem{Rᵤ*mₑ2014/μₑᵤ/0.001,ħ2014,𝘤,2RK2014/𝘤/αinv,mₑ2014}()
const Conventional = UnitSystem{Rᵤ*mₑ1990/μₑᵤ/0.001,ħ1990,𝘤,2RK1990/𝘤/αinv,mₑ1990}()
const English = UnitSystem{kB*rankine/slug/ft^2,ħ/slug/ft^2,𝘤/ft,4π,mₑ/slug}()
const EnglishUS = UnitSystem{1000Rᵤ*mₑ/μₑᵤ*rankine/slug/ftUS^2,ħ/slug/ftUS^2,𝘤/ftUS,4π,mₑ/slug}()

# astronomical units

const GMsun,GMearth,GMjupiter = 1.32712442099e20,398600441.8e6,1.26686534e17
const au,LD,day = 149597870.7e3,384402e3,60^2*24
const pc,ly,GG = au*648000/π,365.25𝘤*day,𝘤*ħ/mP^2
const mₛ = GMsun/GG; const Jₛ = mₛ*au^2/day^2; export mₛ,Jₛ,au,day
const IAU = UnitSystem{Rᵤ*mₑ/μₑᵤ/0.001/Jₛ,ħ/day/Jₛ,day*𝘤/au,4π*1e-7*day^2/Jₛ,mₑ/mₛ}()

# aliased & humorous units

const mf = mass(90/lbm,Metric,English); const Jf = mf*(201.168/14day)^2
const FFF = UnitSystem{1000Rᵤ*mₑ/μₑᵤ*rankine/Jf,ħ/14day/Jf,14day*𝘤/201.168,0,mₑ/mf}()
const units, US, SI, MKS, temp = UnitSystem, UnitSystem, SI2019, Metric, temperature
const CGS, CGS2019, CGSm, CGSe, HLU = Gauss, EMU2019, EMU, ESU, LorentzHeaviside

# natural units

const Planck = UnitSystem{1,1,1,1,√(4π*αG)}()
const PlanckGauss = UnitSystem{1,1,1,4π,√αG}()
const Stoney = UnitSystem{1,αinv,1,4π,√(αG*αinv)}()
const Hartree = UnitSystem{1,1,αinv,4π/αinv^2,1}()
const Rydberg = UnitSystem{1,1,2αinv,π/αinv^2,1/2}()
const Schrodinger = UnitSystem{1,1,αinv,4π/αinv^2,√(αG*αinv)}()
const Electronic = UnitSystem{1,αinv,1,4π,1}()
const Natural = UnitSystem{1,1,1,1,1}()
const NaturalGauss = UnitSystem{1,1,1,4π,1}()
const QCD = UnitSystem{1,1,1,1,1/μₚₑ}()
const QCDGauss = UnitSystem{1,1,1,4π,1/μₚₑ}()
const QCDoriginal = UnitSystem{1,1,1,4π/αinv,1/μₚₑ}()

# physical constants

@pure electronmass(U::typeof(Planck),C::Coupling) = sqrt(4π*coupling(C))
@pure electronmass(U::typeof(PlanckGauss),C::Coupling) = sqrt(coupling(C))
@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,√(αG*αinv)},C::Coupling) where {kB,ħ,𝘤,μ₀} = sqrt(coupling(C)/finestructure(C))
@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,1/μₚₑ},C::Coupling) where {kB,ħ,𝘤,μ₀} = 1/protonelectron(C)
@pure permeability(U::UnitSystem{kB,ħ,𝘤,4π/αinv^2},C::Coupling) where {kB,ħ,𝘤} = 4π*finestructure(C)^2
@pure permeability(U::UnitSystem{kB,ħ,𝘤,π/αinv^2},C::Coupling) where {kB,ħ,𝘤} = π*finestructure(C)^2
@pure lightspeed(U::UnitSystem{kB,ħ,αinv},C::Coupling) where {kB,ħ} = 1/finestructure(C)
@pure lightspeed(U::UnitSystem{kB,ħ,2αinv},C::Coupling) where {kB,ħ} = 2/finestructure(C)
@pure planckreduced(U::UnitSystem{kB,αinv},C::Coupling) where kB = 1/finestructure(C)

@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,mₑ},C::Coupling) where {kB,ħ,μ₀} = electronmass(planck(U),C)
@pure electronmass(U::UnitSystem{kB,ħ,100𝘤,μ₀,1000mₑ},C::Coupling) where {kB,ħ,μ₀} = 1000electronmass(SI,C)
@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,mₑ/1000},C::Coupling) where {kB,ħ,μ₀} = electronmass(SI,C)/1000
@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,mₑ2014},C::Coupling) where {kB,ħ,μ₀} = electronmass(planck(U),C)
@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,mₑ1990},C::Coupling) where {kB,ħ,μ₀} = electronmass(planck(U),C)
@pure electronmass(U::UnitSystem{kB,ħ,𝘤/ftUS,μ₀,mₑ/slug},C::Coupling) where {kB,ħ,μ₀} = electronmass(SI,C)/slug
@pure permeability(U::UnitSystem{kB,ħ,𝘤,μ₀},C::Coupling) where {kB,ħ,𝘤} = finestructure(C)*2𝘩/𝘤/𝘦^2
@pure permeability(U::typeof(ESU2019),C::Coupling) = 1e3*permeability(SI,C)/𝘤^2
@pure permeability(U::typeof(EMU2019),C::Coupling) = 1e7*permeability(SI,C)
@pure permeability(U::typeof(CODATA),C::Coupling) = 2RK2014*finestructure(C)/𝘤
@pure permeability(U::typeof(Conventional),C::Coupling) = 2RK1990*finestructure(C)/𝘤

@pure molarmass(U::UnitSystem{1}) = 1
@pure molarmass(U::UnitSystem{kB},C::Coupling=universe(U)) = NA*electronmass(U,C)/electronunit(C)
@pure molarmass(U::UnitSystem{1e7*kB},C::Coupling=universe(U)) = 1000molarmass(SI2019,C)
@pure molarmass(U::UnitSystem{1e3*kB},C::Coupling=universe(U)) = molarmass(SI2019,C)/1000
@pure molarmass(U::UnitSystem{kB}) where kB = molarmass(CGS)/1000
@pure molarmass(U::UnitSystem{boltzmann(MTS)}) = molarmass(CGS)/1e6
@pure molarmass(U::UnitSystem{boltzmann(CGS)}) = molarmass(Natural)
@pure molarmass(U::UnitSystem{boltzmann(FFF)}) = molarmass(Natural)
@pure molarmass(U::UnitSystem{boltzmann(English)},C::Coupling=universe(U)) = 1000molarmass(SI2019,C)
@pure molarmass(U::UnitSystem{boltzmann(EnglishUS)}) = molarmass(Natural)
@pure molarmass(U::UnitSystem{boltzmann(IAU)}) = 1/1000mₛ

@pure luminousefficacy(U::UnitSystem{1}) = 1
@pure luminousefficacy(U::UnitSystem) = power(Kcd,SI2019,U)

include("kinematic.jl")
include("electromagnetic.jl")
include("thermodynamic.jl")
include("physics.jl")
include("systems.jl")

end # module


module RecursiveFactorization

include("./lu.jl")

let
    while true
        lu!(rand(2,2))
        break
    end
end

end # module


@testset "Coordinates" begin

    @testset "Create" begin
        @testset "Brain Vision" begin
            bv = BrainVision(0, 0, 0)
            @testset "Show" begin
                 show(bv)
            end
        end
        @testset "Talairach" begin
            tal = Talairach(68.3, -26.9, 8.3)
            @testset "Show" begin
                 show(tal)
            end
        end
        @testset "SPM" begin
            spm = SPM(68.3, -26.9, 8.3)
            @testset "Show" begin
                 show(spm)
            end
        end
        @testset "Unknown" begin
            uk = UnknownCoordinate(1.2, 2, 3.2)
            @testset "Show" begin
                 show(uk)
            end
        end
    end

    @testset "Convert" begin
        @testset "MNI -> Talairach" begin

            # Values from table IV in Lancaster et al 2007
            # TODO Find better points to translate

            mni = SPM(73.7, -26.0, 7.0)
            tal = Talairach(68.3, -26.9, 8.3)
            @test isapprox(euclidean(convert(Talairach, mni), tal), 0; atol=2)

            # As an electrode

            # Test as an electrode
            e = Electrode("test", SPM(73.7, -26.0, 7.0), Dict())
            e = conv_spm_mni2tal(e)

            @test isapprox(e.coordinate.x, tal.x; atol = 1.5)
            @test isapprox(e.coordinate.y, tal.y; atol = 1.5)
            @test isapprox(e.coordinate.z, tal.z; atol = 1.5)
            @test isa(e, EEG.Sensor) == true
            @test isa(e, EEG.Electrode) == true
            @test isa(e.coordinate, EEG.Talairach) == true

        end

        @testset "BrainVision -> Talairach" begin

            # TODO this just tests it runs, need to check values

            bv = BrainVision(0, 0, 0)
            tal = Talairach(128, 128, 128)
            @test isapprox(euclidean(convert(Talairach, bv), tal), 0; atol=1.5)

        end
    end

    @testset "Distances" begin

        @test euclidean(Talairach(0, 0, 0), Talairach(1, 1, 1)) == sqrt.(3)

        v = [0, 0, 0]
        @test euclidean(Talairach(1, 1, 1), v) == sqrt.(3)
        @test euclidean(v, Talairach(1, 1, 1)) == sqrt.(3)

    end
end


"""
  module JCRHF
The module required for computation of the wave function using the *Restricted
Hartree-Fock* (RHF) method in a Self-Consistent Field (SCF) calculation. This
module will be used often, as the RHF wave function is often the zeroth-order
wave function for closed-shell systems.
"""
module Energy

using JuliaChem.JCModules
using JuliaChem.JERI

using MPI
using JSON

Base.include(@__MODULE__,"EnergyHelpers.jl")
Base.include(@__MODULE__,"SCF.jl")

"""
  run(input_info::Dict{String,Dict{String,Any}}, basis::Basis)

Execute the JuliaChem RHF algorithm.

One input variable is required:
1. input_info = Information gathered from the input file.
2. basis = The basis set shells, determined from the input file.

One variable is output:
1. scf = Data saved from the SCF calculation.

Thus, proper use of the RHF.run() function would look like this:

```
scf = RHF.run(input_info, basis)
```
"""
function run(mol::Molecule, basis::Basis,
  scf_flags = Dict(); output=0)

  comm=MPI.COMM_WORLD

  if MPI.Comm_rank(comm) == 0 && output >= 2
    println("--------------------------------------------------------------------------------")
    println("                       ========================================                 ")
    println("                                RESTRICTED CLOSED-SHELL                         ")
    println("                                  HARTREE-FOCK ENERGY                           ")
    println("                       ========================================                 ")
    println("")
  end

  #== actually perform scf calculation ==#
  rhfenergy = rhf_energy(mol, basis, scf_flags; output=output)

  if MPI.Comm_rank(comm) == 0 && output >= 2
    println("                       ========================================                 ")
    println("                              END RESTRICTED CLOSED-SHELL                       ")
    println("                                  HARTREE-FOCK ENERGY                           ")
    println("                       ========================================                 ")
    println("--------------------------------------------------------------------------------")
  end

  return rhfenergy
end
export run

end


# * Neon

NeonHF = PP"""! [He] 2s2 2p6
!  Q=8., MEFIT, HF, Ref 22.
ECP,Ne,2,3,0;
1; 2,1.000000,0.00000000;
2; 2,32.029855,112.52822460; 2,12.294136,28.44589604;
2; 2,21.525406,-10.93982872; 2,13.091176,3.39426049;
1; 2,0.850385,-0.16409678;
! References:
! [22] A. Nicklass, M. Dolg, H. Stoll, H. Preuss, J. Chem. Phys. 102, 8942 (1995)."""

NeonWB = PP"""! [He] 2s2 2p6
!  Q=8., MEFIT, WB, Ref 22.
ECP,Ne,2,3,0;
1; 2,1.000000,0.00000000;
2; 2,31.860162,112.52543566; 2,12.362219,28.30083454;
2; 2,21.508034,-11.12658543; 2,12.910447,3.38754919;
1; 2,0.850385,-0.18408921;
! References:
! [22] A. Nicklass, M. Dolg, H. Stoll, H. Preuss, J. Chem. Phys. 102, 8942 (1995)."""

# * Argon

ArgonHF = PP"""! [Ne] 3s2 3p6
!  Q=8., MEFIT, HF, Ref 22.
ECP,Ar,10,4,0;
1; 2,1.000000,0.00000000;
2; 2,10.291970,68.69327818; 2,3.947937,24.48814203;
2; 2,5.388907,27.59672144; 2,2.704463,4.13654627;
2; 2,8.084705,-8.14025335; 2,4.018469,-1.66241708;
1; 2,5.295301,-3.40922653;
! References:
! [22] A. Nicklass, M. Dolg, H. Stoll, H. Preuss, J. Chem. Phys. 102, 8942 (1995)."""

ArgonWB = PP"""! [Ne] 3s2 3p6
!  Q=8., MEFIT, WB, Ref 22.
ECP,Ar,10,4,0;
1; 2,1.000000,0.00000000;
2; 2,10.261721,68.66778801; 2,3.952725,24.04276629;
2; 2,5.392714,27.73076331; 2,2.699967,4.04545904;
2; 2,8.086235,-8.13747696; 2,4.016632,-1.66452808;
1; 2,5.208459,-3.40009845;
! References:
! [22] A. Nicklass, M. Dolg, H. Stoll, H. Preuss, J. Chem. Phys. 102, 8942 (1995)."""

# * Calcium

CalciumDF = PP"""! [Ne] 3s2 3p6 4s2
!  Q=10., MEFIT, MCDHF+Breit, Ref 41.
ECP,Ca,10,4,3;
1; 2,1.,0.;
2; 2,10.556058402,138.832728745; 2,5.293154907,16.754824599;
4; 2,12.714571445,27.661607611; 2,12.878516057,55.297862576; 2,4.147923622,4.210974307; 2,3.841287457,7.627482586;
4; 2,4.703264136,-0.514324716; 2,4.707975898,-0.765184799; 2,13.696151450,-6.490621584; 2,13.698165108,-9.735555558;
2; 2,14.820019137,-14.000148237; 2,22.000803239,-61.616066332;
4; 2,12.714571445,-55.323215221; 2,12.878516057,55.297862576; 2,4.147923622,-8.421948613; 2,3.841287457,7.627482586;
4; 2,4.703264136,0.514324716; 2,4.707975898,-0.510123200; 2,13.696151450,6.490621584; 2,13.698165108,-6.490370372;
2; 2,14.820019137,9.333432158; 2,22.000803239,-30.808033166;

! References:
! [41] I.S. Lim, H. Stoll, P. Schwerdtfeger, J. Chem. Phys. 124, 034107 (2006).
"""sso

# * Krypton

KryptonDF = PP"""! [Ne] 3s2 3p6 3d10 4s2 4p6
!  Q=26., MEFIT, MCDHF+Breit, Ref 36.
ECP,Kr,10,4,3;
1; 2,1.000000,0.000000;
3; 2,70.034410,49.953835; 2,31.895971,369.978238; 2,7.353728,10.054544;
4; 2,47.265183,99.101833; 2,46.664268,198.232524; 2,23.008133,28.204670; 2,22.100979,56.516792;
6; 2,50.768165,-18.588877; 2,50.782228,-27.875398; 2,15.479497,-0.293411; 2,15.559383,-0.469399; 2,2.877154,0.068881; 2,1.991196,0.132270;
2; 2,15.437567,-1.223389; 2,22.055742,-2.991233;
4; 2,47.265183,-198.203666; 2,46.664268,198.232524; 2,23.008133,-56.409339; 2,22.100979,56.516792;
6; 2,50.768165,18.588877; 2,50.782228,-18.583599; 2,15.479497,0.293411; 2,15.559383,-0.312933; 2,2.877154,-0.068881; 2,1.991196,0.088180;
2; 2,15.437567,0.815593; 2,22.055742,-1.495617;
! References:
! [36] K.A. Peterson, D. Figgen, E. Goll, H. Stoll, M. Dolg, J. Chem. Phys. 119, 11113 (2003).
"""sso

KryptonDF2c = PP"""! [Ne] 3s2 3p6 3d10 4s2 4p6
!  Q=26., MEFIT, MCDHF+Breit, Ref 36.
ECP,Kr,10,4,3;
3; 2,70.034410,49.953835; 2,31.895971,369.978238; 2,7.353728,10.054544;
2; 2,46.664268,297.348786; 2,22.100979,84.775187;
3; 2,50.782228,-46.458997; 2,15.559383,-0.782332; 2,1.991196,0.220450;
1; 2,22.055742,-5.234658;
2; 2,47.265183,297.305500; 2,23.008133,84.614009;
3; 2,50.768165,-46.472192; 2,15.479497,-0.733528; 2,2.877154,0.172203;
1; 2,15.437567,-2.854575;
! References:
! [36] K.A. Peterson, D. Figgen, E. Goll, H. Stoll, M. Dolg, J. Chem. Phys. 119, 11113 (2003)."""r

# * Xenon

XenonHF = PP"""! [Ar] 3d10c 4s2c 4p6c 4d10c 5s2 5p6
!  Q=8., MEFIT, HF, Ref 22.
ECP,Xe,46,4,0;
1; 2,1.000000,0.00000000;
2; 2,3.815600,122.76231371; 2,1.878604,8.20135456;
2; 2,3.003078,68.75731963; 2,1.283819,3.64084871;
2; 2,2.027610,23.08319830; 2,0.830435,3.17298823;
2; 2,6.246157,-48.30201724; 2,1.564998,-6.91572892;
! References:
! [22] A. Nicklass, M. Dolg, H. Stoll, H. Preuss, J. Chem. Phys. 102, 8942 (1995)."""

XenonWB = PP"""! [Ar] 3d10c 4s2c 4p6c 4d10c 5s2 5p6
!  Q=8., MEFIT, WB, Ref 22; CPP: alpha=0.8345;delta=0.97391;ncut=1.
ECP,Xe,46,5,0;
1; 2,1.000000,0.00000000;
2; 2,3.940263,122.76382934; 2,2.277264,8.30885115;
2; 2,3.028373,68.82300437; 2,1.394319,3.64674223;
2; 2,2.122605,23.65207854; 2,0.798669,3.25844113;
2; 2,6.164360,-47.70319876; 2,1.542374,-6.54113991;
1; 2,1.847892,-7.10585060;
! References:
! [22] A. Nicklass, M. Dolg, H. Stoll, H. Preuss, J. Chem. Phys. 102, 8942 (1995)."""

XenonDF = PP"""! [Ar] 3d10c 4s2 4p6 4d10 5s2 5p6
!  Q=26., MEFIT, MCDHF+Breit, Ref 36.
ECP,Xe,28,4,3;
1; 2,1.000000,0.000000;
3; 2,40.005184,49.997962; 2,17.812214,281.013303; 2,9.304150,61.538255;
4; 2,15.701772,67.439142; 2,15.258608,134.874711; 2,9.292184,14.663300; 2,8.559003,29.354730;
6; 2,15.185600,35.436908; 2,14.284500,53.195772; 2,7.121889,9.046232; 2,6.991963,13.223681; 2,0.623946,0.084853; 2,0.647284,0.044155;
4; 2,20.881557,-23.089295; 2,20.783443,-30.074475; 2,5.253389,-0.288227; 2,5.361188,-0.386924;
4; 2,15.701772,-134.878283; 2,15.258608,134.874711; 2,9.292184,-29.326600; 2,8.559003,29.354730;
6; 2,15.185600,-35.436908; 2,14.284500,35.463848; 2,7.121889,-9.046232; 2,6.991963,8.815787; 2,0.623946,-0.084853; 2,0.647284,0.029436;
4; 2,20.881557,15.392863; 2,20.783443,-15.037237; 2,5.253389,0.192151; 2,5.361188,-0.193462;
! References:
! [36] K.A. Peterson, D. Figgen, E. Goll, H. Stoll, M. Dolg, J. Chem. Phys. 119, 11113 (2003)."""sso

XenonDF2c = PP"""! [Ar] 3d10c 4s2 4p6 4d10 5s2 5p6
!  Q=26., MEFIT, MCDHF+Breit, Ref 36.
ECP,Xe,28,4,3;
3; 2,40.005184,49.997962; 2,17.812214,281.013303; 2,9.304150,61.538255;
2; 2,15.258608,202.312067; 2,8.559003,44.032095;
3; 2,14.284500,88.659619; 2,6.991963,22.039468; 2,0.647284,0.073591;
2; 2,20.783443,-52.630331; 2,5.361188,-0.677116;
2; 2,15.701772,202.317425; 2,9.292184,43.989900;
3; 2,15.185600,88.592271; 2,7.121889,22.615580; 2,0.623946,0.212133;
2; 2,20.881557,-53.875022; 2,5.253389,-0.672529;
! References:
! [36] K.A. Peterson, D. Figgen, E. Goll, H. Stoll, M. Dolg, J. Chem. Phys. 119, 11113 (2003)."""r

# * Radon

RadonDF = PP"""! [Kr] 4d10c 4f14c 5s2 5p6 5d10 6s2 6p6
!  Q=26., MEFIT, MCDHF+Breit, Ref 36.
ECP,Rn,60,5,4;
1; 2,1.000000,0.000000;
3; 2,30.151242,49.965551; 2,14.521241,283.070000; 2,8.052038,62.002870;
4; 2,11.009942,71.969119; 2,9.617625,143.860559; 2,7.336008,4.714761; 2,6.406253,9.013065;
4; 2,8.369220,36.368365; 2,8.116975,54.551761; 2,5.353656,9.634487; 2,5.097212,14.387902;
4; 2,6.348571,21.797290; 2,6.295949,28.946805; 2,2.882118,1.447365; 2,2.908048,2.177964;
4; 2,11.015205,-25.228095; 2,10.909853,-30.566121; 2,3.482761,-0.574800; 2,3.600418,-0.779538;
4; 2,11.009942,-143.938238; 2,9.617625,143.860559; 2,7.336008,-9.429523; 2,6.406253,9.013065;
4; 2,8.369220,-36.368365; 2,8.116975,36.367841; 2,5.353656,-9.634487; 2,5.097212,9.591935;
4; 2,6.348571,-14.531526; 2,6.295949,14.473402; 2,2.882118,-0.964910; 2,2.908048,1.088982;
4; 2,11.015205,12.614048; 2,10.909853,-12.226448; 2,3.482761,0.287400; 2,3.600418,-0.311815;
! References:
! [36] K.A. Peterson, D. Figgen, E. Goll, H. Stoll, M. Dolg, J. Chem. Phys. 119, 11113 (2003)."""sso


using Visualize, ModernGL
using Visualize.GLRasterization
using FileIO, Images, GeometryTypes

include(Visualize.dir("src", "julia","mesh.jl"))
catmesh = normalmesh(load(Pkg.dir("GLVisualize", "assets", "cat.obj")))
catmesh = Base.view(reinterpret(MeshRender.Vertex, catmesh.parent), catmesh.indexes[1])

proj = perspectiveprojection(42f0, 1f0, 0.1f0, 10f0)
view = lookat(Vec3f0(1.5), Vec3f0(0.5), Vec3f0(0, -1, 0))
solid_color = Vec3f0(0.0, 0.7, 1.0)

light = MeshRender.Light(
    Vec3f0(10),
    Vec3f0(0.0),
    Vec3f0(0.9),
    1.0f0,
    Vec3f0(0.2),
    1.0f0
)
shading = MeshRender.Shading(
    Vec3f0(0.0),
    Vec3f0(0.3),
    8.0f0
)
program_args = (light, shading, (solid_color, proj, view))
events = WindowEvents()
window = events[Window]

uniforms = map(program_args) do x
    UniformBuffer(x)
end;
vbo = VertexArray(catmesh)

draw_cat = GLRasterizer(
    vbo, uniforms,
    MeshRender.vertexshader, MeshRender.fragmentshader
)

glEnable(GL_DEPTH_TEST)
glEnable(GL_CULL_FACE)
glClearColor(1, 1, 1, 0)
GLAbstraction.enabletransparency()

while isopen(window)
    GLWindow.poll_glfw()
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
    draw_cat(vbo, uniforms)
    GLWindow.swapbuffers(window)
end
GLFW.DestroyWindow(window)


export LocalCF, LocalCFMsg, LocalController

const LocalCFMsg = NaiveMsg

"""
A controller framework that only performs local delay compensation without
the knowledge of other agent's controller.
"""
struct LocalCF{N,X,Z,U,S,XDy,ZDy} <: ControllerFramework{X,Z,U,NaiveMsg{X,Z},Nothing}
    central::CentralControl{U,S}
    world_model::WorldDynamics{N,XDy,ZDy}
    delay_model::DelayModel
    "Whether to conpensate communication delay by assuming constant future action."
    conpensate_comm::Bool
end

function LocalCF(
    central::CentralControl{U,S}, world_model::WorldDynamics{N,XDy,ZDy}, delay_model;
    X, Z, conpensate_comm
) where {N,U,S,XDy,ZDy}
    LocalCF{N,X,Z,U,S,XDy,ZDy}(central, world_model, delay_model, conpensate_comm)
end

struct LocalController{id,N,X,Z,U,S,XDy,ZDy} <: Controller{X,Z,U,LocalCFMsg,Nothing}
    cf::LocalCF{N,X,Z,U,S,XDy,ZDy}
    "Central control state."
    s_c::Ref{S}
    τ::Ref{𝕋}
    u::Ref{U}                           # t = τ-1
    u_history::TimedQueue{U}            # t ∈ [τ-Tx,τ+Tu-1]
    t0::𝕋
end

function LocalController(
    id,cf::LocalCF{N,X,Z,U,S,XDy,ZDy}, s_c, t, u0, u_history
) where {N,X,Z,U,S,XDy,ZDy}
    LocalController{id,N,X,Z,U,S,XDy,ZDy}(cf, Ref(s_c), Ref(t), Ref(u0), u_history, t)
end

function OneVision.control!(
    ctrl::LocalController{id,N,X,Z,U},
    x::X,
    z::Z,
    msgs::Each{LocalCFMsg{X,Z}},
)::Tuple{U,Each{LocalCFMsg{X,Z}}} where {id,N,X,Z,U}
    @unpack central, world_model, delay_model, conpensate_comm = ctrl.cf
    @unpack Tx, Tu, Tc, ΔT = short_delay_names(delay_model)
    x_dy, z_dy = world_model.dynamics[id], world_model.obs_dynamics[id]
    τ = ctrl.τ[]
    xt, zt = Timed.(τ-Tx, (x,z))

    push_iter(head, tail) = ((i == 0 ? head : tail[i]) for i in 0:length(tail)-1)

    x_history = self_estimate(x_dy,xt,ctrl.u_history.queue)
    z_history = self_z_estimate(z_dy, zt, push_iter(xt, x_history))

    x_self, z_self = isempty(ctrl.u_history) ? (xt, zt) : (x_history[end], z_history[end])
    msg = LocalCFMsg(x_self[τ+Tu], z_self[τ+Tu])
    if mod(τ - ctrl.t0, ΔT) == 0
        msgs[id] = msg
        xs = [m.x for m in msgs]
        zs = [m.z for m in msgs]
        ctrl.u[] = control_one(central, ctrl.s_c[], xs, zs, τ + Tu, id)
    end
    u = ctrl.u[]
    if conpensate_comm
        x_future = self_estimate(x_dy, x_self, Timed(τ + Tu + t - 1, u) for t in 1:Tc)
        z_future = self_z_estimate(z_dy, z_self, push_iter(x_self, x_future))
        t_rec = τ + Tu + Tc
        msg = LocalCFMsg(x_future[end][t_rec], z_future[end][t_rec])
    end
    msgs′ = fill(msg, N)
    pushpop!(ctrl.u_history, Timed(τ + Tu, u))
    ctrl.τ[] += 1
    u, msgs′
end

function OneVision.make_controllers(
    cf::LocalCF{N,X,Z,U,S},
    init_status::Each{Tuple{X,Z,U}},
    t0::𝕋,
)::Tuple where {N,X,Z,U,S}
    @unpack Tx, Tu, ΔT = short_delay_names(cf.delay_model)
    function mk_controller(id)
        x0, z0, u0 = init_status[id]
        u_history = TimedQueue(u0, t0-Tx, t0+Tu-1)
        ideal_s = init_state(cf.central, t0)

        LocalController(id, cf, ideal_s, t0, u0, u_history)
    end
    function mk_q(_)
        receives = [LocalCFMsg(x, z) for (x, z, u) in init_status]
        constant_queue(receives, msg_queue_length(cf.delay_model))
    end

    ctrls = ntuple(mk_controller, N)
    msg_qs = ntuple(mk_q, N)
    ctrls, msg_qs
end

# LocalCF has nothing to log
OneVision.write_logs(::LocalController) = Dict{𝕋,Nothing}()


export VPGPolicy

"""
Vanilla Policy Gradient

VPGPolicy(;kwargs)

# Keyword arguments
- `approximator`,
- `baseline`,
- `dist`, distribution function of the action
- `γ`, discount factor
- `α_θ`, step size of policy parameter
- `α_w`, step size of baseline parameter
- `batch_size`,
- `rng`,
- `loss`,
- `baseline_loss`,

if the action space is continuous,
then the env should transform the action value, (such as using tanh),
in order to make sure low ≤ value ≤ high
"""
Base.@kwdef mutable struct VPGPolicy{
    A<:NeuralNetworkApproximator,
    B<:Union{NeuralNetworkApproximator,Nothing},
    S,
    R<:AbstractRNG,
} <: AbstractPolicy
    approximator::A
    baseline::B = nothing
    action_space::S
    dist::Any
    γ::Float32 = 0.99f0 # discount factor
    α_θ = 1.0f0 # step size of policy
    α_w = 1.0f0 # step size of baseline
    batch_size::Int = 1024
    rng::R = Random.GLOBAL_RNG
    loss::Float32 = 0.0f0
    baseline_loss::Float32 = 0.0f0
end

"""
About continuous action space, see
* [Diagonal Gaussian Policies](https://spinningup.openai.com/en/latest/spinningup/rl_intro.html#stochastic-policies
* [Clipped Action Policy Gradient](https://arxiv.org/pdf/1802.07564.pdf)
"""

function (π::VPGPolicy)(env::AbstractEnv)
    to_dev(x) = send_to_device(device(π.approximator), x)

    logits = env |> state |> to_dev |> π.approximator |> send_to_host

    if π.action_space isa AbstractVector
        dist = logits |> softmax |> π.dist
        action = π.action_space[rand(π.rng, dist)]
    elseif π.action_space isa Interval
        dist = π.dist.(logits...)
        action = rand.(π.rng, dist)[1]
    else
        error("not implemented")
    end
    action
end

function (π::VPGPolicy)(env::MultiThreadEnv)
    error("not implemented")
    # TODO: can PG support multi env? PG only get updated at the end of an episode.
end

function RLBase.update!(
    trajectory::ElasticSARTTrajectory,
    policy::VPGPolicy,
    env::AbstractEnv,
    ::PreActStage,
    action,
)
    push!(trajectory[:state], state(env))
    push!(trajectory[:action], action)
end

function RLBase.update!(
    t::ElasticSARTTrajectory,
    ::VPGPolicy,
    ::AbstractEnv,
    ::PreEpisodeStage,
)
    empty!(t)
end

RLBase.update!(::VPGPolicy, ::ElasticSARTTrajectory, ::AbstractEnv, ::PreActStage) = nothing

function RLBase.update!(
    π::VPGPolicy,
    traj::ElasticSARTTrajectory,
    ::AbstractEnv,
    ::PostEpisodeStage,
)
    model = π.approximator
    to_dev(x) = send_to_device(device(model), x)

    states = traj[:state]
    actions = traj[:action] |> Array # need to convert ElasticArray to Array, or code will fail on gpu. `log_prob[CartesianIndex.(A, 1:length(A))`
    gains = traj[:reward] |> x -> discount_rewards(x, π.γ)

    for idx in Iterators.partition(shuffle(1:length(traj[:terminal])), π.batch_size)
        S = select_last_dim(states, idx) |> Array |> to_dev
        A = actions[idx]
        G = gains[idx] |> x -> Flux.unsqueeze(x, 1) |> to_dev
        # gains is a 1 column array, but the output of flux model is 1 row, n_batch columns array. so unsqueeze it.

        if π.baseline isa NeuralNetworkApproximator
            gs = gradient(Flux.params(π.baseline)) do
                δ = G - π.baseline(S)
                loss = mean(δ .^ 2) * π.α_w # mse
                ignore() do
                    π.baseline_loss = loss
                end
                loss
            end
            update!(π.baseline, gs)
        elseif π.baseline isa Nothing
            # Normalization. See
            # (http://rail.eecs.berkeley.edu/deeprlcourse-fa17/f17docs/hw2_final.pdf)
            # (https://web.stanford.edu/class/cs234/assignment3/solution.pdf)
            # normalise should not be used with baseline. or the loss of the policy will be too small.
            δ = G |> x -> normalise(x; dims = 2)
        end

        gs = gradient(Flux.params(model)) do
            if π.action_space isa AbstractVector
                log_prob = S |> model |> logsoftmax
                log_probₐ = log_prob[CartesianIndex.(A, 1:length(A))]
            elseif π.action_space isa Interval
                dist = π.dist.(model(S)...) # TODO: this part does not work on GPU. See: https://github.com/JuliaStats/Distributions.jl/issues/1183 .
                log_probₐ = logpdf.(dist, A)
            end
            loss = -mean(log_probₐ .* δ) * π.α_θ
            ignore() do
                π.loss = loss
            end
            loss
        end
        update!(model, gs)
    end
end


export
    MetaAgentCBS

# CBS submodule
module MetaAgentCBS

using ..CRCBS
using Parameters, LightGraphs, DataStructures

@with_kw struct State{S}
    states::Vector{S} = Vector{S}()
end
State(states::T) where {T<:Tuple} = State([states...])
Base.string(s::State) = string("(",prod(map(s->"$(string(s)),",s.states)),")")
@with_kw struct Action{A}
    actions::Vector{A} = Vector{A}()
end
Action(actions::T) where {T<:Tuple} = Action([actions...])
Base.hash(s::State{S}) where {S} = hash(s.states)
Base.:(==)(s1::S,s2::S) where {S<:State} = hash(s1) == hash(s2)
Base.string(a::Action) = string("(",prod(map(a->"$(string(a)),",a.actions)),")")

get_start_index(path::Path{MS,A,C}) where {S<:AbstractGraphState,MS<:State{S},A,C} = get_t(get_initial_state(path).states[1])

abstract type AbstractMetaEnv{S,A,T,C} <: AbstractLowLevelEnv{State{S},Action{A},MetaCostModel{T,C}} end
@with_kw struct LowLevelEnv{S,A,T,C<:AbstractCostModel{T},E<:AbstractLowLevelEnv{S,A,C}} <: AbstractMetaEnv{S,A,T,C}
    envs::Vector{E}             = Vector{E}()
    agent_idxs::Vector{Int}     = Vector{Int}()
    cost_model::MetaCostModel{T,C} = MetaCostModel(
        FullCostModel(sum,TravelTime()),length(envs))
end
# CRCBS.state_type(env::AbstractMetaEnv{S,A,T,C}) where {S,A,T,C} = State{S}
function CRCBS.compute_heuristic_cost(env::MetaAgentCBS.AbstractMetaEnv,cost,sp) where {E<:AbstractLowLevelEnv}
# function compute_heuristic_cost(m::MetaCostModel,::AbstractCostModel,env::E,cost,sp) where {E<:AbstractLowLevelEnv}
    m = get_cost_model(env)
    # h_cost = get_heuristic_cost(env,sp)
    costs = map(i->CRCBS.compute_heuristic_cost(
        get_envs(env)[i],
        cost.independent_costs[i],
        sp.states[i],
        # h_cost[i],
        ),1:m.num_agents)
    CRCBS.MetaCost(costs, CRCBS.aggregate_costs_meta(m.model, costs))
end

"""
    TeamMetaEnv{S,A,T,C,E}

Represents an env where a team of agents must move in rigid formation. The
function `CRCBS.get_possible_actions` must be overridden for all environments to
reflect this rigidity.
"""
@with_kw struct TeamMetaEnv{S,A,T,C<:AbstractCostModel{T},E<:AbstractLowLevelEnv{S,A,C}} <: AbstractMetaEnv{S,A,T,C}
    envs::Vector{E}             = Vector{E}()
    agent_idxs::Vector{Int}     = Vector{Int}()
    cost_model::MetaCostModel{T,C} = MetaCostModel(
        FullCostModel(sum,TravelTime()),length(envs))
end
get_envs(env::AbstractMetaEnv)                 = env.envs
get_agent_idxs(env::AbstractMetaEnv)           = env.agent_idxs
CRCBS.get_cost_model(env::AbstractMetaEnv)     = env.cost_model

function construct_meta_env(envs::Vector{E},idxs::Vector{Int}) where {S,A,T,C<:AbstractCostModel{T},E <: AbstractLowLevelEnv{S,A,C}}
    LowLevelEnv{S,A,T,C,E}(envs=envs,agent_idxs=idxs)
end
function construct_meta_env(envs::Vector{E},idxs::Vector{Int},c_model::C) where {S,A,T,C<:AbstractCostModel{T},E <: AbstractLowLevelEnv{S,A,C}}
    model = MetaCostModel(c_model,length(envs))
    LowLevelEnv{S,A,T,C,E}(envs=envs,agent_idxs=idxs,cost_model=model)
end
function construct_team_env(envs::Vector{E},idxs::Vector{Int}) where {S,A,T,C<:AbstractCostModel{T},E <: AbstractLowLevelEnv{S,A,C}}
    TeamMetaEnv{S,A,T,C,E}(envs=envs,agent_idxs=idxs)
end
function construct_team_env(envs::Vector{E},idxs::Vector{Int},c_model::C) where {S,A,T,C<:AbstractCostModel{T},E <: AbstractLowLevelEnv{S,A,C}}
    model = MetaCostModel(c_model,length(envs))
    TeamMetaEnv{S,A,T,C,E}(envs=envs,agent_idxs=idxs,cost_model=model)
end

function split_meta_path_node(path_node::PathNode{State{S},Action{A}}) where {S,A,C}
    return map(args->PathNode(args...), zip(
        get_s(path_node).states,
        get_a(path_node).actions,
        get_sp(path_node).states
        ))
end
function construct_meta_path_node(nodes::Vector{N}) where {N<:PathNode}
    PathNode(
        State(map(get_s,nodes)),
        Action(map(get_a,nodes)),
        State(map(get_sp,nodes)),
    )
end
construct_meta_path_node(nodes::Tuple) = construct_meta_path_node([nodes...])

# # TODO implement MetaPath
# @with_kw struct MetaPath{S,A,C} <: AbstractPath
#     paths       ::Vector{Path{S,A,C}}   = Vector{Path{S,A,C}}()
#     s0          ::State                 = get_s(get(path_nodes, 1, PathNode{S,A}()))
#     cost        ::MetaCost{C}           = MetaCost(map(get_cost,paths),typemax(C))
# end
# state_type(p::MetaPath{S,A,C}) where {S,A,C}    = State{S}
# action_type(p::MetaPath{S,A,C}) where {S,A,C}   = Action{A}
# cost_type(p::MetaPath{S,A,C}) where {S,A,C}     = MetaCost{C}
# Base.get(p::MetaPath,i,default=node_type(p)()) = construct_meta_path_node(map(p->get(p,i),p.paths))
# Base.getindex(p::MetaPath,i)    = construct_meta_path_node(map(p->getindex(p,i),p.paths))
# function Base.cat(p::P,x::N,args...) where {P<:MetaPath,N<:PathNode}
#     new_paths = typeof(p.paths)()
#     for (path,node) in zip(p.paths,split_path_node(x))
#         setindex!(path,node,i)
#         push!(new_paths,cat(path,node,args...))
#     end
#     P(s0=p.s0,paths=new_paths,cost=p.cost)
# end
# function Base.setindex!(p::P,x,i) where {P<:MetaPath}
#     for (path,node) in zip(p.paths,split_path_node(x))
#         setindex!(path,node,i)
#     end
# end
# Base.length(p::MetaPath) = maximum(map(length,p.paths))
# function Base.push!(p::MetaPath{S,A,C},n::PathNode{State{S},Action{A}}) where {S,A,C}
#     for (path,node) in zip(p.paths,split_path_node(n))
#         push!(path,node)
#     end
#     p
# end
# get_cost(p::MetaPath) = p.cost
# function set_cost!(p::MetaPath,meta_cost::MetaCost)
#     for (path,cost) in zip(p.paths,meta_cost.independent_costs)
#         set_cost!(path)
#     end
#     p.cost = meta_cost.cost
# end
# Base.copy(p::MetaPath) = MetaPath(s0=p.s0,paths=map(copy,p.paths),cost=p.cost)
# get_initial_state(path::MetaPath) = State(map(get_initial_state,path.paths))
# get_final_state(path::MetaPath) = State(map(get_final_state,path.paths))

function construct_meta_path(paths::Vector{P},cost) where {S,A,C,P<:Path{S,A,C}}
    @assert all(map(length,paths) .- maximum(map(length,paths)) .== 0) "$(map(length,paths))"
    starts = map(get_initial_state,paths)
    costs = map(get_cost, paths)
    path_nodes = Vector{PathNode{State{S},Action{A}}}()
    for nodes in zip(map(p->p.path_nodes,paths)...)
        push!(path_nodes, construct_meta_path_node(nodes))
    end
    meta_path = Path(
        path_nodes = path_nodes,
        s0=State(starts),
        cost = MetaCost(costs,cost)
    )
end
function CRCBS.build_env(solver::MetaAgentCBS_Solver, mapf, node, i)
    group = node.solution.group_idxs[i]
    construct_meta_env(
        [build_env(mapf, node, j) for j in group],
        group,
        get_cost_model(mapf.env)
        )
end
CRCBS.get_start(mapf::MAPF,env::AbstractMetaEnv,i) = State([get_start(mapf,e,j) for (e,j) in zip(get_envs(env), get_agent_idxs(env))])
################################################################################
################################################################################
################################################################################

"""
    `split_path(path::Path{State{S},Action{A}}) where {S,A}`

    Helper for breaking a meta-agent path into a set of single-agent paths
"""
function split_path(path::Path{State{S},Action{A},MetaCost{C}}) where {S,A,C}
    N = length(get_s(get_path_node(path, 1)).states)
    paths = [Path{S,A,C}(s0=path.s0.states[i],cost=get_cost(path).independent_costs[i]) for i in 1:N]
    for t in 1:length(path)
        path_node = get_path_node(path, t)
        for (p,n) in zip(paths,split_meta_path_node(path_node))
            push!(p,n)
        end
    end
    return paths
end
function CRCBS.set_solution_path!(solution::MetaSolution, meta_path, group_idx)
    for (idx,path) in zip(solution.group_idxs[group_idx],split_path(meta_path))
        set_solution_path!(solution.solution, path, idx)
    end
end
function CRCBS.set_path_cost!(solution::MetaSolution, meta_cost, group_idx)
    for (idx,cost) in zip(solution.group_idxs[group_idx],meta_cost.independent_costs)
        set_path_cost!(solution.solution, cost, idx)
    end
end
function CRCBS.get_heuristic_cost(env::E,state::State) where {E<:AbstractMetaEnv}
    c = [get_heuristic_cost(e,s) for (e,s) in zip(get_envs(env), state.states)]
end
function CRCBS.states_match(s1::State,s2::State)
    for (state1, state2) in zip(s1.states,s2.states)
        if !states_match(state1,state2)
            return false
        end
    end
    return true
end
function CRCBS.is_goal(env::E,state::State) where {E<:AbstractMetaEnv}
    for (e,s) in zip(get_envs(env), state.states)
        if !is_goal(e,s)
            return false
        end
    end
    return true
end
CRCBS.wait(s::State) = Action(map(CRCBS.wait, s.states))
CRCBS.wait(env::AbstractMetaEnv,s::State) = Action(map(args->CRCBS.wait(args...),zip(get_envs(env),s.states)))
# get_possible_actions
struct ActionIter{A}
    action_lists::Vector{Vector{A}}
end
const ActionIterState = Vector{Int}
function Base.iterate(it::ActionIter)
    iter_state = ActionIterState(ones(Int, length(it.action_lists)))
    if length(iter_state) > 0
        iter_state[end] = 0
        return iterate(it, iter_state)
    else
        return nothing
    end
end
function Base.iterate(iter::ActionIter,iter_state::ActionIterState)
    iter_state = deepcopy(iter_state)
    i = length(iter_state)
    while i > 0
        if iter_state[i] < length(iter.action_lists[i])
            iter_state[i] += 1
            action = Action([iter.action_lists[j][idx] for (j,idx) in enumerate(iter_state)])
            return action, iter_state
        else
            iter_state[i] = 1
            i -= 1
        end
    end
    return nothing
end
function CRCBS.get_possible_actions(env::E,s::State) where {E<:LowLevelEnv}
    return ActionIter([collect(get_possible_actions(env_,s_)) for (env_,s_) in zip(get_envs(env), s.states)])
end
Base.length(iter::ActionIter) = product([length(a) for a in iter.action_lists])
CRCBS.get_next_state(s::State,a::Action) = State(
    [get_next_state(s_,a_) for (s_,a_) in zip(s.states,a.actions)]
    )
CRCBS.get_next_state(env::E,s::State,a::Action) where {E<:AbstractMetaEnv} = State(
    [get_next_state(e_,s_,a_) for (e_,s_,a_) in zip(get_envs(env),s.states,a.actions)]
    )
CRCBS.get_transition_cost(env::E,s::State,a::Action,sp::State) where {E<:AbstractMetaEnv} = [
    get_transition_cost(e_,s_,a_,sp_) for (e_,s_,a_,sp_) in zip(get_envs(env),s.states,a.actions,sp.states)
]
function CRCBS.violates_constraints(env::E, s::State, a::Action, sp::State) where {E<:AbstractMetaEnv}
    for (i, (e_,s_,a_,sp_)) in enumerate(zip(get_envs(env),s.states,a.actions,sp.states))
        if violates_constraints(e_,s_,a_,sp_)
            return true
        end
        # check for self conflict
        P1 = PathNode(s_,a_,sp_)
        for (j, (s2_,a2_,sp2_)) in enumerate(zip(s.states,a.actions,sp.states))
            if i != j
                P2 = PathNode(s2_,a2_,sp2_)
                if detect_state_conflict(P1,P2) || detect_action_conflict(P1,P2)
                    return true
                end
            end
        end
    end
    return false
end

### Reservation Table Interface
function CRCBS.reserve!(table::PIBTReservationTable,env::MetaAgentCBS.AbstractMetaEnv,s,a,sp,t=-1)
    valid = true
    for args in zip(get_envs(env),s.states,a.actions,sp.states)
        valid &= reserve!(table,args...)
    end
    return valid
end
function CRCBS.is_reserved(table::PIBTReservationTable,env::MetaAgentCBS.AbstractMetaEnv,s,a,sp,t=-1)
    for args in zip(get_envs(env),s.states,a.actions,sp.states)
        if is_reserved(table,args...)
            return true
        end
    end
    return false
end
function CRCBS.create_reservations(env::AbstractMetaEnv,s,a,sp,t=-1)
    vcat(map(args->create_reservations(args...,t), zip(get_envs(env),s.states,a.actions,sp.states))...)
end

end


# ==========================================================================================================
# Backend funcion

formula(trm::TableRegressionModel) = trm.mf.f

fe_intercept(f::FormulaTerm) = fe_intercept(f.rhs)
fe_intercept(term::StatsModels.TupleTerm) = map(fe_intercept, term)
fe_intercept(term::FunctionTerm) = first(term.exorig.args) == :fe
fe_intercept(term) = false

width(term::FunctionTerm) = first(term.exorig.args) == :fe ? 0 : 1

# Variable dispersion
dof(model::FixedEffectModel) = model.nobs - model.dof_residual + 1


export K2_5KrylovParams

"""
Type to use the K2.5 formulation with a Krylov method, using the package
[`Krylov.jl`](https://github.com/JuliaSmoothOptimizers/Krylov.jl).
The outer constructor

    K2_5KrylovParams(; uplo = :L, kmethod = :minres, preconditioner = :Identity,
                     rhs_scale = true,
                     atol0 = 1.0e-4, rtol0 = 1.0e-4,
                     atol_min = 1.0e-10, rtol_min = 1.0e-10,
                     ρ0 = sqrt(eps()) * 1e5, δ0 = sqrt(eps()) * 1e5,
                     ρ_min = 1e2 * sqrt(eps()), δ_min = 1e2 * sqrt(eps()),
                     mem = 20)

creates a [`RipQP.SolverParams`](@ref) that should be used to create a [`RipQP.InputConfig`](@ref).
The available methods are:
- `:minres`
- `:minres_qlp`
- `:symmlq`

The list of available preconditioners for this solver is displayed here: [`RipQP.PreconditionerDataK2`](@ref)
"""
mutable struct K2_5KrylovParams <: AugmentedParams
  uplo::Symbol
  kmethod::Symbol
  preconditioner::Symbol
  rhs_scale::Bool
  atol0::Float64
  rtol0::Float64
  atol_min::Float64
  rtol_min::Float64
  ρ0::Float64
  δ0::Float64
  ρ_min::Float64
  δ_min::Float64
  mem::Int
end

function K2_5KrylovParams(;
  uplo::Symbol = :L,
  kmethod::Symbol = :minres,
  preconditioner::Symbol = :Identity,
  rhs_scale::Bool = true,
  atol0::T = 1.0e-4,
  rtol0::T = 1.0e-4,
  atol_min::T = 1.0e-10,
  rtol_min::T = 1.0e-10,
  ρ0::T = sqrt(eps()) * 1e5,
  δ0::T = sqrt(eps()) * 1e5,
  ρ_min::T = 1e2 * sqrt(eps()),
  δ_min::T = 1e3 * sqrt(eps()),
  mem::Int = 20,
) where {T <: Real}
  return K2_5KrylovParams(
    uplo,
    kmethod,
    preconditioner,
    rhs_scale,
    atol0,
    rtol0,
    atol_min,
    rtol_min,
    ρ0,
    δ0,
    ρ_min,
    δ_min,
    mem,
  )
end

mutable struct PreallocatedDataK2_5Krylov{
  T <: Real,
  S,
  L <: LinearOperator,
  Pr <: PreconditionerDataK2,
  Ksol <: KrylovSolver,
} <: PreallocatedDataAugmentedKrylov{T, S}
  pdat::Pr
  D::S                                  # temporary top-left diagonal
  sqrtX1X2::S # vector to scale K2 to K2.5
  tmp1::S # temporary vector for products
  tmp2::S # temporary vector for products
  rhs::S
  rhs_scale::Bool
  regu::Regularization{T}
  δv::Vector{T}
  K::L # augmented matrix
  KS::Ksol
  atol::T
  rtol::T
  atol_min::T
  rtol_min::T
end

function opK2_5prod!(
  res::AbstractVector{T},
  nvar::Int,
  Q::Union{AbstractMatrix{T}, AbstractLinearOperator{T}},
  D::AbstractVector{T},
  A::Union{AbstractMatrix{T}, AbstractLinearOperator{T}},
  sqrtX1X2::AbstractVector{T},
  tmp1::AbstractVector{T},
  tmp2::AbstractVector{T},
  δv::AbstractVector{T},
  v::AbstractVector{T},
  α::T,
  β::T,
  uplo::Symbol,
) where {T}
  tmp2 .= @views sqrtX1X2 .* v[1:nvar]
  mul!(tmp1, Q, tmp2, -α, zero(T))
  tmp1 .= @views sqrtX1X2 .* tmp1 .+ α .* D .* v[1:nvar]
  if β == zero(T)
    res[1:nvar] .= tmp1
  else
    res[1:nvar] .= @views tmp1 .+ β .* res[1:nvar]
  end
  if uplo == :U
    @views mul!(tmp1, A, v[(nvar + 1):end], α, zero(T))
    res[1:nvar] .+= sqrtX1X2 .* tmp1
    @views mul!(res[(nvar + 1):end], A', tmp2, α, β)
  else
    @views mul!(tmp1, A', v[(nvar + 1):end], α, zero(T))
    res[1:nvar] .+= sqrtX1X2 .* tmp1
    @views mul!(res[(nvar + 1):end], A, tmp2, α, β)
  end
  res[(nvar + 1):end] .+= @views (α * δv[1]) .* v[(nvar + 1):end]
end

function PreallocatedData(
  sp::K2_5KrylovParams,
  fd::QM_FloatData{T},
  id::QM_IntData,
  itd::IterData{T},
  pt::Point{T},
  iconf::InputConfig{Tconf},
) where {T <: Real, Tconf <: Real}

  # init Regularization values
  D = similar(fd.c, id.nvar)
  if iconf.mode == :mono
    regu = Regularization(T(sp.ρ0), T(sp.δ0), T(sp.ρ_min), T(sp.δ_min), :classic)
    D .= -T(1.0e0) / 2
  else
    regu =
      Regularization(T(sp.ρ0), T(sp.δ0), T(sqrt(eps(T)) * 1e0), T(sqrt(eps(T)) * 1e0), :classic)
    D .= -T(1.0e-2)
  end
  sqrtX1X2 = fill!(similar(D), one(T))
  tmp1 = similar(D)
  tmp2 = similar(D)
  δv = [regu.δ] # put it in a Vector so that we can modify it without modifying opK2prod!
  K = LinearOperator(
    T,
    id.nvar + id.ncon,
    id.nvar + id.ncon,
    true,
    true,
    (res, v, α, β) ->
      opK2_5prod!(res, id.nvar, fd.Q, D, fd.A, sqrtX1X2, tmp1, tmp2, δv, v, α, β, fd.uplo),
  )

  rhs = similar(fd.c, id.nvar + id.ncon)

  KS = init_Ksolver(K, rhs, sp)

  pdat = eval(sp.preconditioner)(id, fd, regu, D, K)

  return PreallocatedDataK2_5Krylov(
    pdat,
    D,
    sqrtX1X2,
    tmp1,
    tmp2,
    rhs,
    sp.rhs_scale,
    regu,
    δv,
    K, #K
    KS, #K_fact
    T(sp.atol0),
    T(sp.rtol0),
    T(sp.atol_min),
    T(sp.rtol_min),
  )
end

function solver!(
  dd::AbstractVector{T},
  pad::PreallocatedDataK2_5Krylov{T},
  dda::DescentDirectionAllocs{T},
  pt::Point{T},
  itd::IterData{T},
  fd::Abstract_QM_FloatData{T},
  id::QM_IntData,
  res::AbstractResiduals{T},
  cnts::Counters,
  T0::DataType,
  step::Symbol,
) where {T <: Real}

  # erase dda.Δxy_aff only for affine predictor step with PC method
  pad.rhs[1:(id.nvar)] .= @views dd[1:(id.nvar)] .* pad.sqrtX1X2
  pad.rhs[(id.nvar + 1):end] .= @views dd[(id.nvar + 1):end]
  if pad.rhs_scale
    rhsNorm = kscale!(pad.rhs)
  end
  pad.K.nprod = 0
  ksolve!(pad.KS, pad.K, pad.rhs, pad.pdat.P, verbose = 0, atol = pad.atol, rtol = pad.rtol)
  update_kresiduals_history!(res, pad.K, pad.KS.x, pad.rhs)
  if pad.rhs_scale
    kunscale!(pad.KS.x, rhsNorm)
  end
  pad.KS.x[1:(id.nvar)] .*= pad.sqrtX1X2

  dd .= pad.KS.x

  return 0
end

function update_pad!(
  pad::PreallocatedDataK2_5Krylov{T},
  dda::DescentDirectionAllocs{T},
  pt::Point{T},
  itd::IterData{T},
  fd::Abstract_QM_FloatData{T},
  id::QM_IntData,
  res::AbstractResiduals{T},
  cnts::Counters,
  T0::DataType,
) where {T <: Real}
  if cnts.k != 0
    update_regu!(pad.regu)
  end

  if pad.atol > pad.atol_min
    pad.atol /= 10
  end
  if pad.rtol > pad.rtol_min
    pad.rtol /= 10
  end

  # K2.5
  pad.sqrtX1X2 .= one(T)
  pad.sqrtX1X2[id.ilow] .*= sqrt.(itd.x_m_lvar)
  pad.sqrtX1X2[id.iupp] .*= sqrt.(itd.uvar_m_x)
  pad.D .= zero(T)
  pad.D[id.ilow] .-= pt.s_l
  pad.D[id.iupp] .*= itd.uvar_m_x
  pad.tmp1 .= zero(T)
  pad.tmp1[id.iupp] .-= pt.s_u
  pad.tmp1[id.ilow] .*= itd.x_m_lvar
  pad.D .+= pad.tmp1 .- pad.regu.ρ

  pad.δv[1] = pad.regu.δ

  update_preconditioner!(pad.pdat, pad, itd, pt, id, fd, cnts)

  return 0
end


import IntervalRootFinding: BBLeaf, BBNode, BBTree
import IntervalRootFinding: root, nnodes, data, newid, discard_leaf!

import IntervalRootFinding: BreadthFirstBBSearch
import IntervalRootFinding: process, bisect

@testset "Branch and bound tree" begin
    #= Manually create a simple tree with the following structure

        (1)
         |
        (2)
    =#
    leaf = BBLeaf("I'm a leaf", 1, :final)
    node = BBNode(0, [2])
    tree = BBTree(Dict(1 => node),
                  Dict(2 => leaf),
                  [2])

    # Check that printing does not error
    io = IOBuffer()
    println(io, leaf)
    println(io, node)
    println(io, tree)

    @test root(tree) == node
    @test nnodes(tree) == 2

    d = data(tree)
    @test length(d) == 1
    @test d[1] == "I'm a leaf"

    k = newid(tree)
    @test k ∉ 0:2  # The new ID must be new and not 0
    @test isa(k, Integer)

    discard_leaf!(tree, 2)
    # The last leaf was deleted so the tree should be empty
    @test nnodes(tree) == 0

    #= Tests on a slighty more complicated tree

          (1)
         /   \
       (2)   (8)
      /  \     \
    (3)  (5)   (9)
     |   | \     \
    (4) (6) (7)  (10)
    =#

    n1 = BBNode(0, [2, 8])
    n2 = BBNode(1, [3, 5])
    n3 = BBNode(2, [4])
    n4 = BBLeaf("Leaf 4", 3, :working)
    n5 = BBNode(2, [6, 7])
    n6 = BBLeaf("Leaf 6", 5, :final)
    n7 = BBLeaf("Leaf 7", 5, :working)
    n8 = BBNode(1, [9])
    n9 = BBNode(8, [10])
    n10 = BBLeaf("Leaf 10", 9, :working)

    tree = BBTree(Dict(1 => n1, 2 => n2, 3 => n3, 5 => n5, 8 => n8, 9 => n9),
                  Dict(4 => n4, 6 => n6, 7 => n7, 10 => n10),
                  [3, 4, 8])

    # Check that printing does not error
    io = IOBuffer()
    println(io, tree)

    @test newid(tree) ∉ 0:10
    @test length(data(tree)) == 4

    @test nnodes(tree) == 10

    discard_leaf!(tree, 6)
    @test nnodes(tree) == 9

    discard_leaf!(tree, 4)
    @test nnodes(tree) == 7

    discard_leaf!(tree, 10)
    @test nnodes(tree) == 4

    @test length(data(tree)) == 1
end

# Implement the interface for breadth first in a dummy way
struct DummySearch <: BreadthFirstBBSearch{Symbol} end

process(::DummySearch, s::Symbol) = s, s
bisect(::DummySearch, s::Symbol) = :store, :store

@testset "Interface functions" begin
    #= Build the following tree for testing
                    (1)
                   /   \
                (2)     (5: to_bisect)
               /   \
    (3: to_store)  (4: to_discard)
    =#
    rt = BBNode(0, [2, 5])
    node = BBNode(1, [3, 4])
    to_store = BBLeaf(:store, 2, :working)
    to_discard = BBLeaf(:discard, 2, :working)
    to_bisect = BBLeaf(:bisect, 1, :working)

    tree = BBTree(Dict(1 => rt, 2 => node),
                  Dict(3 => to_store, 4 => to_discard, 5 => to_bisect),
                  [3, 4, 5])

    search = DummySearch()

    # First iteration processes the to_store leaf
    tree, _ = iterate(search, tree)

    @test nnodes(tree) == 5
    @test tree[3].status == :final
    @test length(tree.working_leaves) == 2

    # Second iteration processes the to_discard leaf
    tree, _ = iterate(search, tree)

    @test nnodes(tree) == 4
    @test length(tree.working_leaves) == 1

    # Second iteration processes the to_bisect leaf
    tree, _ = iterate(search, tree)

    @test nnodes(tree) == 6
    @test length(tree.working_leaves) == 2
    @test isa(tree[5], BBNode)
end


function readbin(filename)

    a = reinterpret(Float64,read(filename));
    return a

end