raw"""
All the inputs are consumed and the difference of the output sums goes to the
miner

Input hashes and output index point to the original transaction
The unlocking script must match the locking script

Transaction:
  - Version
  - Input Counter
  - Inputs
    - hash
        The transaction origin id
    - output index
        The transacion output idx
    - unlocking script size
    - unlocking script
    - sequence number
  - Outputs
    - amount
    - locking script size
    - locking script

Hexdump of the genesis block area blk00000.dat:

f9be b4d9 1d01 0000 0100 0000 0000 0000  ................
^^^^^^^^^ magic bytes
          ^^^^^^^^^ block size
                    ^^^^^^^^^ version
                              ^^^^^^^^^
0000 0000 0000 0000 0000 0000 0000 0000  ................
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0000 0000 0000 0000 0000 0000 3ba3 edfd  ............;...
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ previous hash
                              ^^^^^^^^^
7a7b 12b2 7ac7 2c3e 6776 8f61 7fc8 1bc3  z{..z.,>gv.a....
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
888a 5132 3a9f b8aa 4b1e 5e4a 29ab 5f49  ..Q2:...K.^J)._I
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ merkle root
                              ^^^^^^^^^ time stamp
ffff 001d 1dac 2b7c 0101 0000 0001 0000  ......+|........
^^^^^^^^^ difficulty
          ^^^^^^^^^ nonce
                    ^^ number of transactions (varint)
                      ^^^^^^^^^^ version
                                ^^ number of inputs
                                   ^^^^
0000 0000 0000 0000 0000 0000 0000 0000  ................
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0000 0000 0000 0000 0000 0000 0000 ffff  ................
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ originating transaction hash
                                   ^^^^
ffff 4d04 ffff 001d 0104 4554 6865 2054  ..M.......EThe T
^^^^ output index of the originating transaction
     ^^ script length (varint)
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
696d 6573 2030 332f 4a61 6e2f 3230 3039  imes 03/Jan/2009
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2043 6861 6e63 656c 6c6f 7220 6f6e 2062   Chancellor on b
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
7269 6e6b 206f 6620 7365 636f 6e64 2062  rink of second b
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
6169 6c6f 7574 2066 6f72 2062 616e 6b73  ailout for banks
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ unlocking script
ffff ffff 0100 f205 2a01 0000 0043 4104  ........*....CA.
^^^^^^^^^ sequence number
          ^^ number of outputs
            ^^^^^^^^^^^^^^^^^^^^ amount (Big endian Float64)
                                ^^ script length
                                   ^^^^
678a fdb0 fe55 4827 1967 f1a6 7130 b710  g....UH'.g..q0..
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
5cd6 a828 e039 09a6 7962 e0ea 1f61 deb6  \..(.9..yb...a..
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
49f6 bc3f 4cef 38c4 f355 04e5 1ec1 12de  I..?L.8..U......
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
5c38 4df7 ba0b 8d57 8a4c 702b 6bf1 1d5f  \8M....W.Lp+k.._
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ac00 0000 00f9 beb4 d9d7 0000 0001 0000  ................
^^ locking script
  ^^^^^^^^^^ lock time
            ^^^^^^^^^^ magic bytes
006f e28c 0ab6 f1b3 72c1 a6a2 46ae 63f7  .o......r...F.c.
4f93 1e83 65e1 5a08 9c68 d619 0000 0000  O...e.Z..h......
0098 2051 fd1e 4ba7 44bb be68 0e1f ee14  .. Q..K.D..h....
677b a1a3 c354 0bf7 b1cd b606 e857 233e  g{...T.......W#>
0e61 bc66 49ff ff00 1d01 e362 9901 0100  .a.fI......b....
0000 0100 0000 0000 0000 0000 0000 0000  ................
0000 0000 0000 0000 0000 0000 0000 0000  ................
0000 00ff ffff ff07 04ff ff00 1d01 04ff  ................
ffff ff01 00f2 052a 0100 0000 4341 0496  .......*....CA..
"""

"""
    TransactionInput

Data Structure Storing Transaction Inputs
"""
struct TransactionInput
    # TODO: remove unlocking_script_size
    # the hash of the transaction where the output comes from
    hash                  :: UInt256
    output_index          :: UInt32
    unlocking_script_size :: UInt64
    unlocking_script      :: Vector{UInt8}
    sequence_number       :: UInt32
end

"""
    TransactionOutput

Data Structure Storing Transaction Outputs
"""
struct TransactionOutput
    # TODO: remove locking_script_size
    amount              :: UInt64 # in satoshis = 1e-8 bitcoins
    locking_script_size :: UInt64
    locking_script      :: Vector{UInt8}
end

"""
    Witness

Data Structure Storing The Witness for SegWit transactions
"""
struct Witness
    data :: Vector{Vector{UInt8}}
end

"""
    Transaction

Data Structure Storing Transactions

To get the hash of a transaction

```julia
double_sha256(tx)
```
"""
struct Transaction
    # TODO: remove *_counter
    version         :: UInt32
    marker          :: UInt8
    flag            :: UInt8
    input_counter   :: UInt64
    inputs          :: Vector{TransactionInput}
    output_counter  :: UInt64
    outputs         :: Vector{TransactionOutput}
    witness_counter :: UInt64
    witnesses       :: Vector{Witness}
    lock_time       :: UInt32
end

function TransactionInput(io)

    # TODO: Cannot read UInt256 directly from io here:
    in_hash = read!(io, Array{UInt256}(undef, 1))[1]
    # in_hash = read(io, UInt256)
    output_index = read(io, UInt32)
    unlocking_script_size = signed(read_varint(io))

    unlocking_script = read!(io, Array{UInt8}(undef, unlocking_script_size))
    sequence_number = read(io, UInt32)

    TransactionInput(
        in_hash,
        output_index,
        unlocking_script_size,
        unlocking_script,
        sequence_number
    )
end

function TransactionOutput(io)

    out_amount = read(io, UInt64)
    out_script_size = signed(read_varint(io))
    locking_script = read!(io, Array{UInt8}(undef, out_script_size))

    TransactionOutput(out_amount,
                      out_script_size,
                      locking_script)
end

function Witness(io::IO)
    n_items = read_varint(io)
    data = Vector{Vector{UInt8}}(undef, n_items)

    for i in 1:n_items
        l = read_varint(io)
        data[i] = read!(io, Array{UInt8, 1}(undef, l))
    end

    Witness(data)
end

function Transaction(io)

    version = read(io, UInt32)

    marker_or_n_in = read_varint(io)

    is_segwit = marker_or_n_in == zero(marker_or_n_in)
    if is_segwit
        marker = UInt8(marker_or_n_in)
        @assert marker == 0x00
        flag = read(io, UInt8)
        @assert flag == 0x01
        n_in = read_varint(io)
    else
        # these two don't matter
        marker = 0xff
        flag = 0xff
        n_in = marker_or_n_in
    end

    @assert n_in > 0
    inputs = TransactionInput[
        TransactionInput(io)
        for i in 1:n_in
    ]

    n_out = read_varint(io)
    @assert n_out > 0
    outputs = TransactionOutput[
        TransactionOutput(io)
        for i in 1:n_out
    ]

    if is_segwit
        n_witness = n_in
        @assert n_witness > 0
        witnesses = Witness[
            Witness(io)
            for i in 1:n_witness
        ]
    else
        n_witness = zero(UInt64)
        witnesses = Witness[]
    end

    locktime = read(io, UInt32)

    Transaction(version,
                marker,
                flag,
                n_in,
                inputs,
                n_out,
                outputs,
                n_witness,
                witnesses,
                locktime)
end

# TODO: is this the same as bytes2hex?
hexarray(x::Array{UInt8}) = mapreduce(x -> string(x, base = 16), *, x)

function showcompact(io::IO, input::TransactionInput)
    # TODO: do something useful here
    println(io, "")
end

function Base.show(io::IO, input::TransactionInput)
    if !get(io, :compact, false)
        println(io, "Transaction input:")
        println(io, "  Hash:                  " * string(input.hash,            base = 16))
        println(io, "  Output index:          " * string(input.output_index,    base = 10))
        println(io, "  Input Sequence:        " * string(input.sequence_number, base = 10))
    end
end

# function Base.showall(io::IO, input::TransactionInput)
#     println(io, "Transaction input:")
#     println(io, "  Hash:                  " * input.hash)
#     println(io, "  Output index:          " * input.output_index)
#     println(io, "  Unlocking script size: " * input.unlocking_script_size)
#     println(io, "  Unlocking script:      " * hexarray(input.unlocking_script))
#     println(io, "  Input Sequence:        " * input.sequence_number)
# end

function Base.show(io::IO, output::TransactionOutput)
    println(io, "Transaction output: " * string(output.amount, base = 10))
end
# function Base.showall(io::IO, output::TransactionOutput)
#     println(io, "Transaction output: "    * string(output.amount,              base = 10))
#     println(io, "  Locking script size: " * string(output.locking_script_size, base = 10))
#     println(io, "  Locking script:      " * hexarray(output.locking_script))
# end

function showcompact(io::IO, transaction::Transaction)
    # TODO: add id here
    println(io, "Transaction: ")
end

function Base.show(io::IO, transaction::Transaction)
    # TODO: add id here
    if !get(io, :compact, false)
        println(io, "Transaction: ")
        println(io, "  Version:        " * string(transaction.version, base = 10))
        println(io, "  Input counter:  " * string(transaction.input_counter, base = 10))
        println(io, "  Output counter: " * string(transaction.output_counter, base = 10))
        println(io, "  Lock time:      " * string(transaction.lock_time, base = 10))
    end
end

# function Base.showall(io::IO, transaction::Transaction)
#     # TODO: add id here
#     println(io, "Transaction: ")
#     println(io, "  Version:        " * string(transaction.version, base = 10))
#     println(io, "  Input counter:  " * string(transaction.input_counter, base = 10))
#     for i in 1:transaction.input_counter
#         show(transaction.inputs[i])
#     end
#     println(io, "  Output counter: " * string(transaction.output_counter, base = 10))
#     for i in 1:transaction.output_counter
#         show(transaction.outputs[i])
#     end
#     println(io, "  Lock time:      " * string(transaction.lock_time, base = 10))
# end

# TODO: little endian only:
function sha256(tx::Transaction)

    ctx = SHA.SHA256_CTX()

    SHA.update!(ctx, to_byte_tuple(tx.version))
    SHA.update!(ctx, to_varint(tx.input_counter))
    for i in 1:tx.input_counter
        SHA.update!(ctx, to_byte_tuple(tx.inputs[i].hash))
        SHA.update!(ctx, to_byte_tuple(tx.inputs[i].output_index))
        SHA.update!(ctx, to_varint(tx.inputs[i].unlocking_script_size))
        SHA.update!(ctx, tx.inputs[i].unlocking_script)
        SHA.update!(ctx, to_byte_tuple(tx.inputs[i].sequence_number))
    end
    SHA.update!(ctx, to_varint(tx.output_counter))
    for i in 1:tx.output_counter
        SHA.update!(ctx, to_byte_tuple(tx.outputs[i].amount))
        SHA.update!(ctx, to_varint(tx.outputs[i].locking_script_size))
        SHA.update!(ctx, tx.outputs[i].locking_script)
    end
    SHA.update!(ctx, to_byte_tuple(tx.lock_time))

    return SHA.digest!(ctx)
end

function double_sha256(tx::Transaction)::UInt256
    tx |> sha256 |> sha256 |> x -> reinterpret(UInt256, x)[1]
end

# TODO: copyto! is save, right?
# TODO: find an allocation free version for these: copyto!(x, i, to_byte_tuple(y))
"""
mutates data in place, advances idx and returns it.
"""
function dump_txin_data!(data, idx, txin::TransactionInput)
    copyto!(data, idx, to_byte_tuple(txin.hash))
    idx += sizeof(txin.hash)

    copyto!(data, idx, to_byte_tuple(txin.output_index))
    idx += sizeof(txin.output_index)

    unlocking_script_size = to_varint(txin.unlocking_script_size)
    copyto!(data, idx, unlocking_script_size)
    idx += sizeof(unlocking_script_size)

    copyto!(data, idx, txin.unlocking_script)
    idx += sizeof(txin.unlocking_script)

    copyto!(data, idx, to_byte_tuple(txin.sequence_number))
    idx += sizeof(txin.sequence_number)

    return idx
end

# TODO: copyto! is save, right?
# TODO: find an allocation free version for these: copyto!(x, i, to_byte_tuple(y))
"""
mutates data in place, advances idx and returns it.
"""
function dump_txout_data!(data, idx, txout::TransactionOutput)
    copyto!(data, idx, to_byte_tuple(txout.amount))
    idx += sizeof(txout.amount)

    locking_script_size = to_varint(txout.locking_script_size)
    copyto!(data, idx, locking_script_size)
    idx += sizeof(locking_script_size)

    copyto!(data, idx, txout.locking_script)
    idx += sizeof(txout.locking_script)

    return idx
end

# TODO: copyto! is save, right?
# TODO: find an allocation free version for these: copyto!(x, i, to_byte_tuple(y))
"""
mutates data in place, advances idx and returns it.
"""
function dump_tx_data!(data, idx, tx::Transaction)
    copyto!(data, idx, to_byte_tuple(tx.version))
    idx += sizeof(tx.version)

    is_segwit = tx.marker == 0x00
    if is_segwit
        data[idx] = tx.marker
        idx += 1
        data[idx] = tx.flag
        idx += 1
    end

    input_counter = to_varint(tx.input_counter)
    copyto!(data, idx, input_counter)
    idx += sizeof(input_counter)
    for i in 1:tx.input_counter
        idx = dump_txin_data!(data, idx, tx.inputs[i])
    end

    output_counter = to_varint(tx.output_counter)
    copyto!(data, idx, output_counter)
    idx += sizeof(output_counter)
    for i in 1:tx.output_counter
        idx = dump_txout_data!(data, idx, tx.outputs[i])
    end

    if is_segwit
        for i in 1:tx.input_counter
            n_items = length(tx.witnesses[i].data)
            n_items_bytes = to_varint(n_items)
            copyto!(data, idx, n_items_bytes)
            idx += length(n_items_bytes)

            for j in 1:n_items
                l = length(tx.witnesses[i].data[j])
                l_bytes = to_varint(l)
                copyto!(data, idx, l_bytes)
                idx += length(l_bytes)

                copyto!(data, idx, tx.witnesses[i].data[j])
                idx += l
            end
        end
    end

    copyto!(data, idx, to_byte_tuple(tx.lock_time))
    idx += sizeof(tx.lock_time)

    return idx
end

# TODO: common abstraction with "Link"?
struct FilePointer
    file_number   :: UInt64
    file_position :: UInt64
end

function Transaction(fp::FilePointer)
    fp.file_number |>
        get_block_chain_file_path |>
        x -> open(x) do fh
            seek(fh, fp.file_position)
            Transaction(fh)
        end
end

total_output(tx::Transaction) = sum(x -> x.amount, tx.outputs)

is_coinbase(txin::TransactionInput) = txin.hash == zero(UInt256)


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

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

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

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

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

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

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

typealias Batch Vector{Beam}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


"""
```
DSGEVECM{T} <: AbstractDSGEVECMModel{T}
```
implements a simple interface for combining a given DSGE model
with a VECM (a VAR with cointegrating terms) to create a DSGE-VECM.
Confer with Del Negro and Schorfheide (2006)
and/or Del Negro, Schorfheide, Smets, and Wouters (2007) for details about DSGE-VECMs.
We recommend the first paper as an initial introduction to DSGE-VECMs.

The recommended constructor requires the user to provide
(1) an `AbstractDSGEModel` object, (2) which structural shocks from
the DSGE to use, and (3) the subspec (optional, defaults to "ss0").
If the subspec "ss0" is used, then the result is a `DSGEVECM`
whose VECM component is "empty" in that the
observables, cointegrating relationships, lags, and λ weight are not specified.
The reason why this constructor requires the user to specify which
structural shocks of DSGE to use is that this information is
DSGE-specific rather than information about the VECM.

However, we can also construct a `DSGEVECM` without having to
specify the structural shocks when calling the constructor,
although we still need to give an instance of an `AbstractDSGEModel`.

### Example
The code below instantiates an empty `DSGEVECM`
with `AnSchorfheide` as the underlying DSGE and then
calls `update!` on the empty DSGE-VECM
to add information about the desired DSGE-VECM spec.

```jldoctest; output = false
dsgevecm = DSGEVECM(AnSchorfheide())
DSGE.update!(dsgevecm, shocks = [:rm_sh, :z_sh, :g_sh],
    observables = [:obs_gdp, :obs_cpi, :obs_nominalrate],
    λ = 1., lags = 4)

# output

DSGE-VECM Model
observables:      3
data_vintage:     200310
DSGE model:       an_schorfheide
DSGE description: Julia implementation of model defined in 'Bayesian Estimation of DSGE Models' by Sungbae An and Frank Schorfheide: AnSchorfheide, ss0
```

### Fields

#### DSGE object
* `dsge::AbstractDSGEModel{T}`: underlying DSGE model object

#### DSGE-VECM Information
* `observables::OrderedDict{Symbol,Int}`: dictionary mapping observables
    of the VECM to their index in the matrices representing the DSGE-VECM.
* `cointegrating::OrderedDict{Symbol,Int}`: dictionary mapping cointegrating relationships
    of the VECM to their index in the matrices representing the DSGE-VECM.
    When creating the state space representation of a DSGE-VECM, the
    cointegrating relationships will come after the observables. Accordingly,
    the first `cointegrating` index will be after the last observable.
* `cointegrating_add::OrderedDict{Symbol,Int}`: dictionary mapping
    additional cointegrating relationships of the VECM to their index
    in the matrices representing the DSGE-VECM. These relationships are
    intended to be ones that only add to the `DD` matrix in DSGE.jl's
    representation of a state space model. Importantly, these relationships
    do not require changing the `ZZ` matrix. As a result, the indices for
    these relationships start at 1.
* `shocks::OrderedDict{Symbol,Int}`: dictionary mapping structural
    shocks in the DSGE to their index in the matrices representing the DSGE-VECM
* `lags::Int`: number of lags in the VECM
* `λ::T`: weight on the DSGE prior

#### Auxiliary Information
* `spec::String`: concatenates `dsgevar` with the spec of the DSGE, e.g.
    for `AnSchorfheide`, we have `dsgevar_an_schorfheide`.
* `subspec::String`: specifies the model subspecification.
    Cached here for filepath computation.
* `testing::Bool`: indicates whether the model is in testing mode.
    Currently, this setting has no uses in practice
"""
mutable struct DSGEVECM{T} <: AbstractDSGEVECMModel{T}
    dsge::AbstractDSGEModel{T}
    observables::OrderedDict{Symbol,Int}
    cointegrating::OrderedDict{Symbol,Int}
    cointegrating_add::OrderedDict{Symbol,Int}
    shocks::OrderedDict{Symbol,Int}
    lags::Int
    λ::T
    spec::String
    subspec::String
    testing::Bool
end

function Base.show(io::IO, m::DSGEVECM)
    @printf io "DSGE-VECM Model\n"
    @printf io "observables:      %i\n" n_observables(m)
    @printf io "data_vintage:     %s\n" data_vintage(m.dsge)
    @printf io "DSGE model:       %s\n" spec(get_dsge(m))
    @printf io "DSGE description: %s\n" description(get_dsge(m))
end

function DSGEVECM(dsge::AbstractDSGEModel{T}, shocks::Vector{Symbol}, subspec::String = "ss0";
                 custom_settings::Dict{Symbol, Setting} = Dict{Symbol, Setting}(),
                 copy_dsge::Bool = false, testing = false) where {T<:Real}

    # Initialize specs
    spec     = "dsgevecm_" * ModelConstructors.spec(dsge)
    subspec  = subspec

    # Initialize empty DSGEVECM
    if copy_dsge
        dsge = deepcopy(dsge)
    end
    m = DSGEVECM{T}(dsge, OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(),
                    OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), 0,
                    0., spec, subspec, testing)

    # Initialize shocks from DSGE
    update!(m; shocks = shocks, check_valid = false)

    # Initialize subspec
    init_subspec!(m)

    # Do checks of observables, shocks, and lags
    check_valid_dsgevecm(m)

    return m
end

# Empty constructor
function DSGEVECM(dsge::AbstractDSGEModel{T}, subspec::String = "ss0";
                 custom_settings::Dict{Symbol, Setting} = Dict{Symbol, Setting}(),
                 copy_dsge::Bool = false, testing = false) where {T<:Real}

    # Initialize specs
    spec     = "dsgevar_" * ModelConstructors.spec(dsge)
    subspec  = subspec

    # Initialize empty DSGEVECM
    if copy_dsge
        dsge = deepcopy(dsge)
    end
    m = DSGEVECM{T}(dsge, OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(),
                    OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), 0,
                    0., spec, subspec, testing)

    # Initialize shocks from DSGE
    update!(m; check_valid = false)

    # Initialize subspec
    init_subspec!(m)

    # Do checks of observables, shocks, and lags
    check_valid_dsgevecm(m)

    return m
end


# The simplified model of ECME-RIRR plus ROS generation by ETC ODE system
using Parameters, StaticArrays

include("common.jl")
include("cicr16.jl")
include("ck.jl")
include("force.jl")
include("ik2009.jl")
include("ina.jl")
include("inaca2009.jl")
include("ins_bg.jl")
include("katp.jl")
include("mtdna.jl")
include("oxphos-ode.jl")
include("pump.jl")
include("ros.jl")
include("tca.jl")
include("transmito.jl")
include("cabuf.jl")
include("concentrations.jl")

# Initial conditions (largely from Zhou et al. (2009))
u0_tup = (
    vm = -87.28,  # Sarcolemmal membrane potential (mV)
    dpsi = 175.318,  # Mitochondrial membrane potential (mV)
    m_na = 0.0327,  # Fast Na gating (activation)
    h_na = 0.9909,  # Fast Na gating (inactivation)
    j_na = 0.9941,  # Fast Na gating (slow inactivation)
    x_k = 1.1212e-4,  # K gating (activation)
    na_i = 8.2143,  # Cytoplasmic Na (mM)
    k_i = 150.8,  # Cytoplasmic K (mM)
    ca_i = 5.8653e-5,  # Cytoplasmic Ca (mM)
    ca_jsr = 0.1948,  # Junctional SR Ca (mM)
    ca_nsr = 0.1948,  # Network SR Ca (mM)
    ca_ss = 7.057e-5,  # Subspace calcium
    ltr_ca = 8.949E-3,  # Ca bound to low affinity troponin sites (mM)
    htr_ca = 0.1321,  # Ca bound to high affinity troponin sites (mM)
    # CICR states
    po1 = 8.309e-5,
    po2 = 7.12e-11,
    pc2 = 0.2471,
    c0 = 0.9991,
    c1 = 8.175e-5,
    c2 = 2.508e-9,
    c3 = 3.421e-14,
    c4 = 1.749e-20,
    o = 2.624e-20,
    cca0 = 1.1328e-3,
    cca1 = 1.7591e-8,
    cca2 = 6.3826e-13,
    cca3 = 1.815e-15,
    y_ca = 0.9479,
    # Tropomyosin cross bridge states
    p0 = 2.601e-5,
    p1 = 2.248e-5,
    p2 = 4.199e-5,
    p3 = 3.657e-5,
    n1 = 2.243e-5,
    adp_i = 8 - 7.8707,  # Cytosolic ATP(mM), linked to EC coupling
    adp_ic = 8 - 7.7107,  # Cytosolic ATP(mM), not linked to EC coupling
    crp_i = 5.0297,  # Cytosolic creatine phosphate(mM), linked to EC coupling
    crp_ic = 5.1291,  # Cytosolic creatine phosphate(mM), not linked to EC coupling
    # TCA cycle states
    ca_m = 2.190477831750499e-5,
    adp_m = 0.005788892127721868,
    nadh = 9.267368089681039,
    isoc = 0.4858673722382066,
    akg = 0.0001907868710833752,
    scoa = 0.28625004579156677,
    suc = 0.00011243840322880198,
    fum = 0.005275867239849815,
    mal = 0.002781903362877464,
    oaa = 2.683029248919834e-8,
    # ROS
    sox_m = 0.0004966725362784072,
    sox_i = 8.570060710246488e-9,
    h2o2_i = 2.061879987223193e-9,
    gsh_i = 0.9999899536989296,
    # ETC states
    Q_n = 1.3268400422082767,
    Qdot_n = 0.06757658200311795,
    QH2_n = 0.6002709581315435,
    QH2_p = 0.6001505595245672,
    Qdot_p = 0.07820141736951611,
    Q_p = 1.3269604407629754,
    b1 = 0.11362728970178769,
    b2 = 0.14764414671393097,
    b3 = 0.04917170786399363,
    b4 = 0.014656855720284924,
    fes_ox = 0.054765936912457895,
    cytc1_ox = 0.16423864061514593,
    cytc_ox = 0.1081945884323241)

nameLUT = Dict((k=>i) for (i, k) in enumerate(keys(u0_tup)))

# TODO: conservation of Q pool
Q_POOL = u0_tup[:Q_n] + u0_tup[:Qdot_n] + u0_tup[:QH2_n] + u0_tup[:QH2_p] + u0_tup[:Qdot_p] + u0_tup[:Q_p]
# TODO: b pool = C3 conc
B_POOL = u0_tup[:b1] + u0_tup[:b2] + u0_tup[:b3] + u0_tup[:b4]

conc = Concentrations(nadph_i=1.0,
                      pi_m=2.0,
                      mg_m=0.4,
                      mg_i=3.1,
                      AXP_M_T=1.5,
                      NAD_T=10.0,
                      AXP_I_T=8.0,
                      O2=6E-3)

@with_kw struct Params
    conc = conc
    pSODi = SODParams(K1=1.2E3, K3=24.0, K5=2.5E-4, ET=1.43E-3, KI_H2O2=0.5)
    pGPXi = GPXParams(ET=0.01, 𝚽1=5e-3, 𝚽2=0.75)
    pGRi = GRParams(K1=5E-3, ET=0.01, KM_GSSG=0.06, KM_NADPH=0.015, GSH_T=1.0)
    pCAT = CATParams(K1=17.0, ET=0.01, FR=0.05)
    pC1 = C1Params()
    pC2 = C2Params()
    pC3 = C3Params()
    pC4 = C4Params()
    pC5 = C5Params(ρF1=0.05)
    pANT = ANTParams(VMAX=5E-3)
    pMNCE = MNCEParams(VMAX=7.3E-4)
    pMCU = MCUParams(VMAX=0.028)
    pIMAC = IMACParams()
    pCS = CSParams(ACCOA=1.0, KCAT=0.5, ET=0.4, KM_ACCOA = 0.0126, KM_OAA = 6.4E-4)
    pACO = ACOParams(TCA_T=1.0, KF=1.25E-2, KEQ=2.22)
    pIDH = IDHParams(KM_ADP=0.62, KM_CA=5E-4, KI_NADH=0.19, KCAT=0.05, ET=0.109, KH1=8.1e-5, KH2=5.98E-5, KM_ISOC=1.52, NI_ISOC=2, KM_NAD=0.923)
    pKGDH = KGDHParams(KM_MG=0.0308, KM_CA=1.27E-3, ET=0.5, KCAT=0.075, KM_AKG=1.94, KM_NAD=38.7, NI_AKG=1.2)
    pSL = SLParams(KF=5E-3, KEQ=3.115, COA=0.02)
    pFH = FHParams(KF=3.32E-3, KEQ=1.0)
    pMDH = MDHParams(KH1=1.131E-5, KH2=26.7, KH3=6.68E-9, KH4=5.62E-6, K_OFFSET=3.99E-2, KCAT=0.11, ET=0.15, KM_MAL=1.49, KI_OAA=3.1E-3, KM_NAD=0.22)
    pAAT = AATParams(KF=6.44E-4, KEQ=6.6, KASP=1.5E-6, GLU=10.0)
    pDOX = DOXParams()
    pMTDNA = MTDNAParams()
    # G_H_MITO = 2E-6  # Mitochondrial membrane conductance (Gauthier 2013)
    G_H_MITO = 1E-8  # (Li 2015)
    δCA = 3E-4  # Ca buffering factor
    R_TR= inv(9.09) # Rate of Ca diffusion from NSR to JSR (1/ms)
    R_XFER = inv(0.5747)  # Rate of Ca diffusion from subspace to cytosol (1/ms)
    G_NA = 12.8  # Sodium conductance
    V_NSR = 1.4E-6  # Network SR Volume (μL)
    V_JSR = 0.16E-6  # Junctional SR Volume (μL)
    V_SS = 0.495E-9  # Total subspace Volume (μL)
    V_MYO = 25.84E-6  # Cytoplasm Volume (μL)
    V_MITO = 15.89E-6  # Mitochiondrial Volume (μL)
    A_CAP = 1.534E-4  # Cell capacitance area (cm²)
    A_CAP_V_MYO_F = A_CAP / (V_MYO * F)
    V_MITO_V_MYO = V_MITO / V_MYO
    pIK = IKParams()
    pCK = CKParams()
    pXB = XBParams()
    pLCC = LCCParams()
    pRYR = RYRParams()
    pSERCA = SERCAParams(KMF_CA=2.4e-4, KMR_CA=1.64269, NFB=1.4, NRB=1)
    pPMCA = PMCAParams()
    pNS = NSParams()
    pNKA = NKAParams()
    pNCX = NCXParams()
    pKATP = KATPParams()
    pCSQN = CABUFParams(KM_CA=0.8, ET=5.0)
    pCMDN = CABUFParams(KM_CA=2.38E-3, ET=5E-2)
    pBG = BGParams()
    BCL = 0
end

function _istim(t, bcl=0, peroid=0.5, strength=-80.0)
    if (bcl > 0) && (rem(t, bcl) <= peroid)
        return strength
    else
        return zero(strength)
    end
end

function ecme_dox(u, p, t)
    (vm, dpsi, m_na, h_na, j_na, x_k, na_i, k_i, ca_i, ca_jsr, ca_nsr, ca_ss, ltr_ca, htr_ca,
     po1, po2, pc2, c0, c1, c2, c3, c4, o, cca0, cca1, cca2, cca3, y_ca, p0, p1, p2, p3, n1,
     adp_i, adp_ic, crp_i, crp_ic, ca_m, adp_m, nadh, isoc, akg, scoa, suc, fum, mal, oaa,
     sox_m, sox_i, h2o2_i, gsh_i, Q_n, Qdot_n, QH2_n, QH2_p, Qdot_p, Q_p, b1, b2, b3, b4,
     fes_ox, cytc1_ox, cytc_ox) = u

    @unpack conc = p
	@unpack ca_o, na_o, k_o, AXP_I_T, h_i, h_m, mg_i, mg_m, ΔpH, FAC_PH, O2, nadph_i, pi_m, AXP_I_T = conc
    vfrt = vm * F_RT
    evfrtm1 = expm1(vfrt)
	ΔμH = _ΔμH(dpsi, ΔpH)  # proton motive force
    atp_i = AXP_I_T - adp_i
    atp_ic = AXP_I_T - adp_ic

    # CICR
    @unpack pLCC, pRYR, R_TR, R_XFER = p
	(d_c0, d_c1, d_c2, d_c3, d_c4, d_o, d_cca0, d_cca1, d_cca2,
     d_cca3) = lcc_system(c0, c1, c2, c3, c4, o, cca0, cca1, cca2, cca3, ca_ss, vm, pLCC)
    d_y_ca = dyca(y_ca, vm)
    (d_po1, d_po2, d_pc2) = ryr_system(po1, po2, pc2, ca_ss, pRYR)
    (iCaL, iCaK) = ical(y_ca, o, vfrt, evfrtm1, k_i, k_o, ca_o, pLCC)
    jRel = jrel(po1, po2, ca_jsr, ca_ss, pRYR.R_RYR)
    jTr = R_TR * (ca_nsr - ca_jsr)
    jXfer = R_XFER * (ca_ss - ca_i)
    # Myofibril
    @unpack pCK, pXB = p
    (d_p0, d_p1, d_p2, d_p3, d_n1, d_ltr_ca, d_htr_ca) = xb_system(p0, p1, p2, p3, n1, ltr_ca, htr_ca, ca_i, pXB)
	vAM = vam(p0, p1, p2, atp_i, adp_i, pXB)
    jTrpn = d_ltr_ca + d_htr_ca

    # CK shuttle
    (d_atp_in_ck, d_atp_ic, d_crp_i, d_crp_ic) = ck_system(atp_i, atp_ic, crp_i, crp_ic, pCK, AXP_I_T)

    # Other ion channels
    @unpack G_CAB, G_NAB = p.pBG
    @unpack G_KP, P_NA_K, G_K, G_K1 = p.pIK
    @unpack pNKA, pPMCA, pSERCA, pKATP, G_NA, pNS, pNCX = p
    d_m_na = d_mna(m_na, vm)
    d_h_na = d_hna(h_na, vm)
    d_j_na = d_jna(j_na, vm)
    eNa = _ena(na_o, na_i)
    iNa = ina(vm, eNa, m_na, h_na, j_na, G_NA)
    d_x_k = dxk(x_k, vm)
    iK = ik(vm, _enak(k_o, na_o, k_i, na_i, P_NA_K), x_k, G_K)
	eK = _ek(k_o, k_i)
    iKp = ikp(vm, eK, G_KP)
    iK1 = ik1(vm, eK, G_K1)
	iKatp = ikatp(vm, eK, na_i, atp_i, adp_i, mg_i, pKATP)
    iNsNa = insna(na_i, ca_i, vfrt, evfrtm1, na_o, pNS)
    iNaCa = inaca(na_i, ca_i, vfrt, evfrtm1, ca_o, pNCX)
    iCaB = icab(vm, _eca(ca_o, ca_i), G_CAB)
    iNaB = inab(vm, eNa, G_NAB)
    iNaK = inak(na_i, atp_i, adp_i, vm, evfrtm1, pNKA)
    iPCa = ipca(ca_i, atp_i, adp_i, pPMCA)
    jUp = jup(ca_i, ca_nsr, atp_i, adp_i, pSERCA)

    # Ion fluxes across mitochondria
    @unpack pMCU, G_H_MITO, pMNCE = p
    vUni = vuni(ca_i, dpsi, pMCU)
    vNaCa = vnaca(ca_i, ca_m, na_i, dpsi, pMNCE, pMCU)
    vHleak = vhleak(ΔμH, G_H_MITO)

    # ETC
    @unpack pC1, pC2, pC3, pC4, pC5, pDOX, pANT = p
    @unpack C2_INHIB, C3_INHIB, C4_INHIB = pDOX
    @unpack MT_PROT = p.pMTDNA
    vANT = vant(adp_m, atp_i, dpsi, conc, pANT)
    vc1, vROSC1 = c1_rates(nadh, Q_n, QH2_n, dpsi, h_m, sox_m, conc, pC1, pDOX, MT_PROT)
    vc2 = vSDH = vsdh_mm(Q_n, QH2_n, suc, fum, oaa, pC2, C2_INHIB)
    vE, vHresC4, vO2 = c4_rates(cytc_ox, dpsi, h_m, conc, pC4, C4_INHIB, MT_PROT)
    (vATPase, vHu) = c5_rates(ΔμH, pi_m, adp_m, conc, pC5, MT_PROT)
    # Complex III rates
    @unpack (KD_Q, K03, KEQ3, α, δ₁, K04, KEQ4_OX, KEQ4_RD, β, δ₂, KEQ6, K06, γ, δ₃,
        KEQ7_OX, K07_OX, KEQ7_RD, K07_RD, KEQ8_OX, K08_OX, KEQ8_RD, K08_RD, K09,
        KEQ9, K010, KEQ10, K33, KEQ33, ρC3) = pC3
    C3_CONC = MT_PROT * ρC3
	C4_CONC = MT_PROT * pC4.ρC4
    fes_rd = C3_CONC - fes_ox
    cytc1_rd = C3_CONC - cytc1_ox
	cytc_rd = C4_CONC - cytc_ox
	MYXOTHIAZOLE = ANTIMYCIN = 1

	# v1 = Q reduction
	v1 = vc1 + vc2
	# v2 = QH2 diffusion (n-side -> p-side)
    v2 = KD_Q * (QH2_n - QH2_p)
    # v3 = QH2 to FeS
	v3 = _v3(fes_ox, QH2_p, fes_rd, Qdot_p, K03, KEQ3, FAC_PH, MYXOTHIAZOLE)
	# v4 = Qdot_p to bH
	el = _el(α, δ₁, dpsi)
	er = _er(α, δ₁, dpsi)
	v4_ox = _v4(b1, Qdot_p, el, b2, Q_p, er, K04, KEQ4_OX)
	v4_rd = _v4(b3, Qdot_p, el, b4, Q_p, er, K04, KEQ4_RD)
    # v5 = Q diffusion (p-side -> n-side)
    v5 = KD_Q * (Q_p - Q_n)
    # v6 = bH to bL
	v6 = _v6(b2, b3, dpsi, K06, KEQ6, β, δ₂)
    # v7 = bL to Qn; v8: bL to Qdot_n
	el = _el(γ, δ₃, dpsi)
	er = _er(γ, δ₃, dpsi)
	v7_ox = _v7(b3, Q_n, el, b1, Qdot_n, er, K07_OX, KEQ7_OX, ANTIMYCIN) * C3_INHIB
	v7_rd = _v7(b4, Q_n, el, b2, Qdot_n, er, K07_RD, KEQ7_RD, ANTIMYCIN) * C3_INHIB
	# v8 = bL to Qdot_n and proton transfer
	v8_ox = _v8(b3, Qdot_n, el, b1, QH2_n, er, K08_OX, KEQ8_OX, FAC_PH, ANTIMYCIN) * C3_INHIB
	v8_rd = _v8(b4, Qdot_n, el, b2, QH2_n, er, K08_RD, KEQ8_RD, FAC_PH, ANTIMYCIN) * C3_INHIB
    # v9 = fes -> cytc1
    v9 = _v9(fes_rd, cytc1_ox, fes_ox, cytc1_rd, K09, KEQ9)
    # v10 = Qdot + O2 -> O2- + Q  (ROS produced by complex III)
    vROSC3 = v10 = _v10(O2, Qdot_p, sox_m, Q_p, K010, KEQ10)

	# Redox reaction between cytc1 and cytc
	v33 = _v3(cytc1_rd, cytc_ox, cytc1_ox, cytc_rd, K33, KEQ33)

    # ODE system of complex III
    d_Q_n = v5 - v7_ox - v7_rd - v1
    d_Qdot_n = v7_ox + v7_rd - v8_ox - v8_rd
    d_QH2_n = v8_ox + v8_rd - v2 + v1
    d_QH2_p = v2 - v3
    d_Qdot_p = v3 - v10 - v4_ox - v4_rd
    d_Q_p = v10 + v4_ox + v4_rd - v5
    d_b1 = v7_ox + v8_ox - v4_ox
    d_b2 = v4_ox + v7_rd + v8_rd - v6
    d_b3 = v6 - v4_rd - v7_ox - v8_ox
    d_b4 = v4_rd - v7_rd - v8_rd
    d_fes_ox = v9 - v3
    d_cytc1_ox = v33 - v9
	d_cytc_ox = vE - v33

	vROS = vROSC1 + vROSC3
    vHres = 4 * vc1 + 2 * v3 + vHresC4

    # TCA cycle rates
    @unpack pSL, pCS, pACO, pIDH, pKGDH, pFH, pMDH, pAAT = p
    nad = _nad(nadh, conc)
	vCS  = vcs(oaa, pCS)
	vACO = vaco(isoc, akg, scoa, suc, fum, mal, oaa, pACO)
	vIDH = vidh(isoc, nadh, nad, adp_m, ca_m, h_m, mg_m, pIDH)
	vKGDH = vkgdh(akg, nadh, nad, ca_m, mg_m, pKGDH)
	vSL = vsl(scoa, suc, adp_m, _atp_m(adp_m, conc), pSL)
	vFH = vfh(fum, mal, pFH)
	vMDH = vmdh(mal, oaa, nadh, nad, h_m, pMDH)
	vAAT = vaat(akg, oaa, pAAT)

    # ODE system of TCA cycle
	d_nadh = -vc1 + vIDH + vKGDH + vMDH
	d_isoc = vACO - vIDH
	d_akg = vIDH + vAAT - vKGDH
	d_scoa = vKGDH - vSL
	d_suc = vSL - vSDH
	d_fum = vSDH - vFH
	d_mal = vFH - vMDH
	d_oaa = vMDH - vCS - vAAT
    d_adp_m = vANT - vSL - vATPase

    @unpack (pSODi, pGPXi, pGRi, pCAT, pIMAC, V_MITO_V_MYO) = p
    vSOD_i = vsod(sox_i, h2o2_i, pSODi)
    vGPX_i = vgpx(h2o2_i, gsh_i, pGPXi)
    vGR_i = vgr(gsh_i, nadph_i, pGRi)
    vCAT = v_cat(h2o2_i, pCAT)
    vIMAC, vTrROS = vimac_vtrros(dpsi, sox_m, sox_i, pIMAC)

    # ODE system of ROS scavenging
    d_sox_m = vROS - vTrROS
    d_sox_i = V_MITO_V_MYO * vTrROS - vSOD_i
	d_h2o2_i = 0.5 * vSOD_i - vGPX_i - vCAT
	d_gsh_i  = vGR_i - vGPX_i

	# Remaining ODE equations
	@unpack (δCA, A_CAP_V_MYO_F, V_MYO, V_SS, V_NSR, V_JSR, V_MITO_V_MYO, BCL) = p
	iStim = _istim(t, BCL)
	d_vm = -CM_INV * (iNa + iCaL + iK + iK1 + iKp + iNaCa + iNaK + iNsNa + iPCa + iCaB + iNaB + iStim + iKatp + iCaK)
	d_na_i = -A_CAP_V_MYO_F * (iNa + iNaB + iNsNa + 3 * (iNaCa + iNaK)) # + V_MITO_V_MYO * (vNHE - 3 * vNaCa)
    d_k_i = -A_CAP_V_MYO_F * (iK + iK1 + iKp + iKatp + iStim - 2 * iNaK + iCaK)
	d_atp_i = d_atp_in_ck + V_MITO_V_MYO * vANT - vAM - 0.5 * jUp - A_CAP_V_MYO_F * (iPCa + iNaK)

	@unpack pCSQN, pCMDN, A_CAP = p
	d_ca_i = _β_ca(ca_i, pCMDN) * (jXfer - jUp - jTrpn - 0.5 * A_CAP_V_MYO_F * (iPCa + iCaB - 2 * iNaCa) + V_MITO_V_MYO * (vNaCa - vUni))
	d_ca_nsr = _β_ca(ca_nsr, pCSQN) * (V_MYO * jUp - V_JSR * jTr) / V_NSR
    d_ca_jsr = _β_ca(ca_jsr, pCSQN) * (jTr - jRel)
    d_ca_ss = _β_ca(ca_ss, pCMDN) * (V_JSR * jRel - V_MYO * jXfer - 0.5 * A_CAP / F * iCaL) / V_SS
    d_ca_m = δCA * (vUni - vNaCa)

    d_dpsi = CM_MITO_INV * (vHres - vHu - vANT - vHleak - vNaCa - 2 * vUni - vIMAC)
    # @show t, na_i, m_na, h_na, j_na, d_m_na, d_h_na, d_j_na, iNa

	# Disable rates for debugging
    #=
	d_vm = 0
	d_m_na = d_h_na = d_j_na = d_na_i = 0
	d_x_k = d_k_i = 0
	d_p0 = d_p1 = d_p2 = d_p3 = d_n1 = 0
	d_ca_i = d_ca_jsr = d_ca_nsr = d_ca_ss = d_ltr_ca = d_htr_ca = 0
	d_atp_i = d_atp_ic = d_crp_i = d_crp_ic = 0
    d_po1 = d_po2 = d_pc2 = d_c0 = d_c1 = d_c2 = d_c3 = d_c4 = d_o = d_cca0 = d_cca1 = d_cca2 = d_cca3 = d_y_ca = 0
    =#

    d_adp_i = -d_atp_i
    d_adp_ic = -d_atp_ic
    du = @SVector[d_vm, d_dpsi, d_m_na, d_h_na, d_j_na, d_x_k, d_na_i, d_k_i, d_ca_i, d_ca_jsr, d_ca_nsr, d_ca_ss, d_ltr_ca, d_htr_ca,
     d_po1, d_po2, d_pc2, d_c0, d_c1, d_c2, d_c3, d_c4, d_o, d_cca0, d_cca1, d_cca2, d_cca3, d_y_ca, d_p0, d_p1, d_p2, d_p3, d_n1,
     d_adp_i, d_adp_ic, d_crp_i, d_crp_ic, d_ca_m, d_adp_m, d_nadh, d_isoc, d_akg, d_scoa, d_suc, d_fum, d_mal, d_oaa,
     d_sox_m, d_sox_i, d_h2o2_i, d_gsh_i, d_Q_n, d_Qdot_n, d_QH2_n, d_QH2_p, d_Qdot_p, d_Q_p, d_b1, d_b2, d_b3, d_b4,
     d_fes_ox, d_cytc1_ox, d_cytc_ox,
     iCaL, iCaK, jRel, jTr, jXfer, vAM, jTrpn, iNa, iK, iKp, iK1, iKatp, iNsNa, iNaCa, iCaB,
     iNaB, iNaK, iPCa, jUp, vUni, vNaCa, vHleak, vANT, vc1, vROSC1, vSDH, vE, vHresC4, vO2,
     vATPase, vHu, vROSC3, vHres, vCS, vACO, vIDH, vKGDH, vSL, vFH, vMDH, vAAT, vSOD_i, vGPX_i,
     vGR_i, vCAT, vIMAC, vTrROS]
end

rateSymbols = [:d_vm, :d_dpsi, :d_m_na, :d_h_na, :d_j_na, :d_x_k, :d_na_i, :d_k_i, :d_ca_i, :d_ca_jsr, :d_ca_nsr, :d_ca_ss, :d_ltr_ca, :d_htr_ca,
 :d_po1, :d_po2, :d_pc2, :d_c0, :d_c1, :d_c2, :d_c3, :d_c4, :d_o, :d_cca0, :d_cca1, :d_cca2, :d_cca3, :d_y_ca, :d_p0, :d_p1, :d_p2, :d_p3, :d_n1,
 :d_adp_i, :d_adp_ic, :d_crp_i, :d_crp_ic, :d_ca_m, :d_adp_m, :d_nadh, :d_isoc, :d_akg, :d_scoa, :d_suc, :d_fum, :d_mal, :d_oaa,
 :d_sox_m, :d_sox_i, :d_h2o2_i, :d_gsh_i, :d_Q_n, :d_Qdot_n, :d_QH2_n, :d_QH2_p, :d_Qdot_p, :d_Q_p, :d_b1, :d_b2, :d_b3, :d_b4,
 :d_fes_ox, :d_cytc1_ox, :d_cytc_ox,
 :iCaL, :iCaK, :jRel, :jTr, :jXfer, :vAM, :jTrpn, :iNa, :iK, :iKp, :iK1, :iKatp, :iNsNa, :iNaCa, :iCaB,
 :iNaB, :iNaK, :iPCa, :jUp, :vUni, :vNaCa, :vHleak, :vANT, :vc1, :vROSC1, :vSDH, :vE, :vHresC4, :vO2,
 :vATPase, :vHu, :vROSC3, :vHres, :vCS, :vACO, :vIDH, :vKGDH, :vSL, :vFH, :vMDH, :vAAT, :vSOD_i, :vGPX_i,
 :vGR_i, :vCAT, :vIMAC, :vTrROS]

rateMap = Dict((k => i) for (i, k) in enumerate(rateSymbols))

# Ignore the intermediate rates
rhs(u, p, t) = ecme_dox(u, p, t)[1:length(u)]

reject_step(u, p, t) = any( x -> (x < 0.0), @view u[3:end])


# All the types supported by mshadow.
typealias DType Union{Float32, Float64, Float16, UInt8, Int32}
@enum TypeFlag kFloat32 kFloat64 kFloat16 kUint8 kInt32
typealias DEFAULT_DTYPE Float32

function toTypeFlag{T <: DType}(:: Type{T})
  if T == Float32
    return kFloat32
  elseif T == Float64
    return kFloat64
  elseif T == Float16
    return kFloat16
  elseif T == UInt8
    return kUint8
  elseif T == Int32
    return kInt32
  else
    throw(ArgumentError("Can't convert $T to DType."))
  end
end

function fromTypeFlag(T :: TypeFlag)
  if T == kFloat32
    return Float32
  elseif T == kFloat64
    return Float64
  elseif T == kFloat16
    return Float16
  elseif T == kUint8
    return UInt8
  elseif T == kInt32
    return Int32
  else
    throw(ArgumentError("Can't convert DType $T."))
  end
end

# create a NDArray handle of specific shape
function _ndarray_alloc{N}(shape :: NTuple{N, Int}, ctx :: Context, delay_alloc :: Bool)
  h_ref  = Ref{MX_handle}(0)
  shape  = flipdim(MX_uint[shape...],1)
  @mxcall(:MXNDArrayCreate, (Ptr{MX_uint}, MX_uint, Cint, Cint, Cint, Ref{MX_handle}),
      shape, length(shape), ctx.device_type, ctx.device_id, delay_alloc, h_ref)
  handle = MX_NDArrayHandle(h_ref[])
  return handle
end

# create a NDArray handle of specific shape type
function _ndarray_alloc{T <: DType,N}(:: Type{T}, shape :: NTuple{N, Int}, ctx :: Context, delay_alloc :: Bool)
  h_ref  = Ref{MX_handle}(0)
  shape  = flipdim(MX_uint[shape...],1)
  dtype  = toTypeFlag(T)
  @mxcall(:MXNDArrayCreateEx, (Ptr{MX_uint}, MX_uint, Cint, Cint, Cint, Cint, Ref{MX_handle}),
      shape, length(shape), ctx.device_type, ctx.device_id, delay_alloc, dtype, h_ref)
  handle = MX_NDArrayHandle(h_ref[])
  return handle
end

# create a handle to an empty NDArray, this handle can be used to hold
# results returned by libmx API calls
function _ndarray_alloc()
  h_ref = Ref{MX_handle}(0)
  @mxcall(:MXNDArrayCreateNone, (Ref{MX_handle},), h_ref)
  return MX_NDArrayHandle(h_ref[])
end

################################################################################
# NDArray Type
################################################################################
"""
    NDArray

Wrapper of the `NDArray` type in `libmxnet`. This is the basic building block
of tensor-based computation.

!!! note
      since C/C++ use row-major ordering for arrays while Julia follows a
      column-major ordering. To keep things consistent, we keep the underlying data
      in their original layout, but use *language-native* convention when we talk
      about shapes. For example, a mini-batch of 100 MNIST images is a tensor of
      C/C++/Python shape (100,1,28,28), while in Julia, the same piece of memory
      have shape (28,28,1,100).
"""
type NDArray
  handle   :: MX_NDArrayHandle
  writable :: Bool

  function NDArray(handle, writable=true)
    new(handle, writable)
  end
end

function Base.show(io :: IO, arr :: NDArray)
  print(io, "mx.NDArray{$(eltype(arr))}$(size(arr))")
end

function NDArray{T<:Real}(data :: Array{T})
  copy(data, cpu())
end

function Base.unsafe_convert(::Type{MX_handle}, obj::NDArray)
  Base.unsafe_convert(MX_handle, obj.handle)
end
Base.convert(t::Type{MX_handle}, obj::NDArray) = Base.unsafe_convert(t, obj)
Base.cconvert(t::Type{MX_handle}, obj::NDArray) = Base.unsafe_convert(t, obj)

################################################################################
# NDArray functions exported to the users
################################################################################
"""
    context(arr :: NDArray)

Get the context that this `NDArray` lives on.
"""
function context(arr :: NDArray)
  ref_typeid = Ref{Cint}(0)
  ref_devid  = Ref{Cint}(0)
  @mxcall(:MXNDArrayGetContext, (MX_handle, Ref{Cint}, Ref{Cint}),
          arr, ref_typeid, ref_devid)
  return Context(ref_typeid[], ref_devid[])
end

"""
    empty(DType, shape :: Tuple, ctx :: Context)
    empty(DType, shape :: Tuple)
    empty(DType, dim1, dim2, ...)

Allocate memory for an uninitialized `NDArray` with a specified type.
"""
function empty{N,T<:DType}(::Type{T}, shape :: NTuple{N, Int})
  empty(T, shape, cpu())
end
function empty{N,T<:DType}(:: Type{T}, shape :: NTuple{N, Int}, ctx :: Context)
  NDArray(_ndarray_alloc(T, shape, ctx, false))
end
function empty{T<:DType}(:: Type{T}, shape :: Int...)
  empty(T, shape)
end

"""
    empty(shape :: Tuple, ctx :: Context)
    empty(shape :: Tuple)
    empty(dim1, dim2, ...)

Allocate memory for an uninitialized `NDArray` with specific shape of type Float32.
"""
function empty{N}(shape :: NTuple{N, Int})
  empty(shape, cpu())
end
function empty{N}(shape :: NTuple{N, Int}, ctx :: Context)
  NDArray(_ndarray_alloc(shape, ctx, false))
end
function empty(shape :: Int...)
  empty(shape)
end

"""
    zeros(DType, shape :: Tuple, ctx :: Context)
    zeros(DType, shape :: Tuple)
    zeros(DType, dim1, dim2, ...)

Create zero-ed `NDArray` with specific shape and type
"""
function zeros{N,T<:DType}(:: Type{T}, shape :: NTuple{N, Int})
  zeros(T, shape, cpu())
end
function zeros{N,T<:DType}(:: Type{T}, shape :: NTuple{N, Int}, ctx :: Context)
  arr = empty(T, shape, ctx)
  arr[:] = zero(T)
  return arr
end
function zeros{T<:DType}(:: Type{T}, shape :: Int...)
  zeros(T, shape)
end

"""
    zeros(shape :: Tuple, ctx :: Context)
    zeros(shape :: Tuple)
    zeros(dim1, dim2, ...)

Create zero-ed `NDArray` with specific shape.
"""
function zeros{N}(shape :: NTuple{N, Int})
  zeros(shape, cpu())
end
function zeros{N}(shape :: NTuple{N, Int}, ctx :: Context)
  arr = empty(shape, ctx)
  arr[:] = 0
  return arr
end
function zeros(shape :: Int...)
  zeros(shape)
end

"""
    ones(DType, shape :: Tuple, ctx :: Context)
    ones(DType, shape :: Tuple)
    ones(DType, dim1, dim2, ...)

Create an `NDArray` with specific shape & type, and initialize with 1.
"""
function ones{N,T<:DType}(:: Type{T}, shape :: NTuple{N, Int})
  ones(T, shape, cpu())
end
function ones{N,T<:DType}(:: Type{T}, shape :: NTuple{N, Int}, ctx :: Context)
  arr = empty(T, shape, ctx)
  arr[:] = one(T)
  return arr
end
function ones{T<:DType}(:: Type{T}, shape :: Int...)
  ones(T, shape)
end

"""
    ones(shape :: Tuple, ctx :: Context)
    ones(shape :: Tuple)
    ones(dim1, dim2, ...)

Create an `NDArray` with specific shape and initialize with 1.
"""
function ones{N}(shape :: NTuple{N, Int})
  ones(shape, cpu())
end
function ones{N}(shape :: NTuple{N, Int}, ctx :: Context)
  arr = empty(shape, ctx)
  arr[:] = 1
  return arr
end
function ones(shape :: Int...)
  ones(shape)
end

import Base: size, length, ndims, eltype

"""
    size(arr :: NDArray)
    size(arr :: NDArray, dim :: Int)

Get the shape of an `NDArray`. The shape is in Julia's column-major convention. See
also the notes on NDArray shapes [`NDArray`](@ref).
"""
function size(arr :: NDArray)
  ref_ndim  = Ref{MX_uint}(0)
  ref_shape = Ref{Ptr{MX_uint}}(0)
  @mxcall(:MXNDArrayGetShape, (MX_handle, Ref{MX_uint}, Ref{Ptr{MX_uint}}),
          arr, ref_ndim, ref_shape)
  tuple(map(Int, flipdim(unsafe_wrap(Array, ref_shape[], ref_ndim[]),1))...)
end
function size(arr :: NDArray, dim :: Int)
  size(arr)[dim]
end

"""
    length(arr :: NDArray)

Get the number of elements in an `NDArray`.
"""
function length(arr :: NDArray)
  prod(size(arr))
end

"""
    ndims(arr :: NDArray)

Get the number of dimensions of an `NDArray`. Is equivalent to `length(size(arr))`.
"""
function ndims(arr :: NDArray)
  length(size(arr))
end

"""
    eltype(arr :: NDArray)

Get the element type of an `NDArray`.
"""
function eltype{T <: Union{NDArray, MX_NDArrayHandle}}(arr :: T)
  dtype_ref = Ref{Cint}(0)
  @mxcall(:MXNDArrayGetDType, (MX_handle, Ptr{Cint}), arr, dtype_ref)

  if dtype_ref[] == -1 # arr->is_none()
    warn("Eltype of $arr is not defined")
    Base.show_backtrace(STDOUT,backtrace())
    println()
    return Float32
  else
    return fromTypeFlag(TypeFlag(dtype_ref[]))
  end
end

import Base: slice
"""
    slice(arr :: NDArray, start:stop)

Create a view into a sub-slice of an `NDArray`. Note only slicing at the slowest
changing dimension is supported. In Julia's column-major perspective, this is the last
dimension. For example, given an `NDArray` of shape (2,3,4), `slice(array, 2:3)` will create
a `NDArray` of shape (2,3,2), sharing the data with the original array. This operation is
used in data parallelization to split mini-batch into sub-batches for different devices.
"""
function slice(arr :: NDArray, ::Colon)
  arr
end
function slice(arr :: NDArray, slice::UnitRange{Int})
  dim1 = size(arr)[end]
  @assert(1 <= slice.start <= slice.stop <= dim1)
  if slice.start == 1 && slice.stop == dim1
    return arr
  end

  hdr_ref = Ref{MX_handle}(0)
  # note Julia is 1-based, inclusive-inclusive indexing, while C++ is
  # 0-based, inclusive-exclusive indexing. So 1:3 in Julia should
  # translates into 0:3 in C++.
  @mxcall(:MXNDArraySlice, (MX_handle, MX_uint, MX_uint, Ref{MX_handle}),
          arr, slice.start-1, slice.stop, hdr_ref)
  return NDArray(MX_NDArrayHandle(hdr_ref[]), arr.writable)
end

import Base: setindex!

"""
    setindex!(arr :: NDArray, val, idx)

Assign values to an `NDArray`. Elementwise assignment is not implemented, only the following
scenarios are supported

* `arr[:] = val`: whole array assignment, `val` could be a scalar or an array (Julia `Array`
  or `NDArray`) of the same shape.
* `arr[start:stop] = val`: assignment to a *slice*, `val` could be a scalar or an array of
  the same shape to the slice. See also [`slice`](@ref).
"""
function setindex!(arr :: NDArray, val :: Real, ::Colon)
  @assert(arr.writable)
  _set_value(convert(eltype(arr), val), arr)
  return arr
end
function setindex!{T<:Real}(arr :: NDArray, val :: Array{T}, ::Colon)
  copy!(arr, val)
end
function setindex!(arr :: NDArray, val :: NDArray, ::Colon)
  copy!(arr, val)
end
function setindex!{T<:Real}(arr :: NDArray, val :: Union{T,Array{T},NDArray}, idx::UnitRange{Int})
  setindex!(slice(arr, idx), val, Colon())
end

import Base: getindex
"""
    getindex(arr :: NDArray, idx)

Shortcut for [`slice`](@ref). A typical use is to write

```julia
  arr[:] += 5
```

which translates into

```julia
  arr[:] = arr[:] + 5
```

which furthur translates into

```julia
  setindex!(getindex(arr, Colon()), 5, Colon())
```

!!! note
    The behavior is quite different from indexing into Julia's `Array`. For example, `arr[2:5]`
    create a **copy** of the sub-array for Julia `Array`, while for `NDArray`, this is
    a *slice* that shares the memory.
"""
function getindex(arr :: NDArray, ::Colon)
  return arr
end

"""
Shortcut for [`slice`](@ref). **NOTE** the behavior for Julia's built-in index slicing is to create a
copy of the sub-array, while here we simply call `slice`, which shares the underlying memory.
"""
function getindex(arr :: NDArray, idx::UnitRange{Int})
  slice(arr, idx)
end

import Base: copy!, copy, convert
"""
    copy!(dst :: Union{NDArray, Array}, src :: Union{NDArray, Array})

Copy contents of `src` into `dst`.
"""
function copy!(dst :: NDArray, src :: NDArray)
  @assert(dst.writable)
  if dst.handle == src.handle
    warn("Copying an NDArray to itself")
    return
  end

  _copyto(src, dst)
  return dst
end

function copy!{T<:DType}(dst :: Array{T}, src :: NDArray)
  @assert T == eltype(src)
  @assert size(dst) == size(src)
  @mxcall(:MXNDArraySyncCopyToCPU, (MX_handle, Ptr{Void}, Csize_t),
          src, pointer(dst), length(dst))
  return dst
end
function copy!{T<:Real}(dst :: Array{T}, src :: NDArray)
  copy!(dst, copy(src))
end

function copy!{T<:Real}(dst :: NDArray, src :: Array{T})
  @assert dst.writable
  @assert size(dst) == size(src)
  src = convert(Array{eltype(dst)}, src) # this might involve copying
  @mxcall(:MXNDArraySyncCopyFromCPU, (MX_handle, Ptr{Void}, Csize_t),
          dst.handle, pointer(src), length(src))
  return dst
end

function copy_ignore_shape!{T<:Real}(dst :: NDArray, src :: Array{T})
  @assert dst.writable
  @assert length(dst) == length(src)
  src = convert(Array{eltype(dst)}, src) # this might involve copying
  @mxcall(:MXNDArraySyncCopyFromCPU, (MX_handle, Ptr{Void}, Csize_t),
          dst.handle, pointer(src), length(src))
  return dst
end

"""
    copy(arr :: NDArray)
    copy(arr :: NDArray, ctx :: Context)
    copy(arr :: Array, ctx :: Context)

Create a copy of an array. When no `Context` is given, create a Julia `Array`.
Otherwise, create an `NDArray` on the specified context.
"""
# Create copy: NDArray -> Julia Array
function copy(arr :: NDArray)
  j_arr = Array{eltype(arr)}(size(arr))
  copy!(j_arr, arr)
end

# Create copy: NDArray -> NDArray in a given context
function copy(arr :: NDArray, ctx :: Context)
  dst = NDArray(_ndarray_alloc(eltype(arr), size(arr), ctx, true))
  copy!(dst, arr)
end

# Create copy: Julia Array -> NDArray in a given context
function copy{T<:DType}(arr :: Array{T}, ctx :: Context)
  dst = empty(T, size(arr), ctx)
  copy!(dst, arr)
end

"""
    convert(::Type{Array{T}}, arr :: NDArray)

Convert an `NDArray` into a Julia `Array` of specific type. Data will be copied.
"""
function convert{T<:Real}(t::Type{Array{T}}, arr :: NDArray)
  convert(t, copy(arr))
end

"""
    @inplace

Julia does not support re-definiton of `+=` operator (like `__iadd__` in python),
When one write `a += b`, it gets translated to `a = a+b`. `a+b` will allocate new
memory for the results, and the newly allocated `NDArray` object is then assigned
back to a, while the original contents in a is discarded. This is very inefficient
when we want to do inplace update.

This macro is a simple utility to implement this behavior. Write

```julia
  @mx.inplace a += b
```

will translate into

```julia
  mx.add_to!(a, b)
```

which will do inplace adding of the contents of `b` into `a`.
"""
macro inplace(stmt)
  if stmt.head == :+= || stmt.head == :.+=
    Expr(:call, :add_to!, esc(stmt.args[1]), esc(stmt.args[2]))
  elseif stmt.head == :-= || stmt.head == :.-=
    Expr(:call, :sub_from!, esc(stmt.args[1]), esc(stmt.args[2]))
  elseif stmt.head == :.*=
    Expr(:call, :mul_to!, esc(stmt.args[1]), esc(stmt.args[2]))
  elseif stmt.head == :./=
    Expr(:call, :div_from!, esc(stmt.args[1]), esc(stmt.args[2]))
  else
    error("unsupported inplace translation for $stmt")
  end
end

"""
    add_to!(dst :: NDArray, args :: Union{Real, NDArray}...)

Add a bunch of arguments into `dst`. Inplace updating.
"""
function add_to!(dst :: NDArray, args :: Union{Real, NDArray}...)
  @assert dst.writable
  for arg in args
    if isa(arg, Real)
      _plus_scalar(dst, convert(eltype(dst), arg), dst)
    else
      _plus(dst, arg, dst)
    end
  end
  return dst
end

import Base: +, .+

"""
    +(args...)
    .+(args...)

Summation. Multiple arguments of either scalar or `NDArray` could be
added together. Note at least the first or second argument needs to be an `NDArray` to
avoid ambiguity of built-in summation.
"""
function +(arg0 :: NDArray, args :: Union{Real, NDArray}...)
  ret = copy(arg0, context(arg0))
  add_to!(ret, args...)
end
function .+(arg0 :: NDArray, args :: Union{Real, NDArray}...)
  +(arg0, args...)
end
function +(arg0 :: Real, arg1 :: NDArray, args :: Union{Real, NDArray}...)
  +(arg1, arg0, args...)
end
function .+(arg0 :: Real, arg1 :: NDArray, args :: Union{Real, NDArray}...)
  .+(arg1, arg0, args...)
end

"""
    sub_from!(dst :: NDArray, args :: Union{Real, NDArray}...)

Subtract a bunch of arguments from `dst`. Inplace updating.
"""
function sub_from!(dst :: NDArray, arg :: Union{Real, NDArray})
  @assert dst.writable
  if isa(arg, Real)
    _minus_scalar(dst, convert(eltype(dst), arg), dst)
  else
    _minus(dst, arg, dst)
  end
end

import Base: -, .-

"""
    -(arg0, arg1)
    -(arg0)
    .-(arg0, arg1)

Subtraction `arg0 - arg1`, of scalar types or `NDArray`. Or create
the negative of `arg0`.
"""
function -(arg0 :: NDArray, arg1 :: Union{Real, NDArray})
  ret = copy(arg0, context(arg0))
  sub_from!(ret, arg1)
end
function .-(arg0 :: NDArray, arg1 :: Union{Real, NDArray})
  -(arg0, arg1)
end
function -(arg0 :: Real, arg1 :: NDArray)
  ret = -arg1
  add_to!(ret, arg0)
  return ret
end
function .-(arg0 :: Real, arg1 :: NDArray)
  -(arg0, arg1)
end

function -(arg0 :: NDArray)
  _mul_scalar(arg0, -one(eltype(arg0)))
end

"""
    mul_to!(dst :: NDArray, arg :: Union{Real, NDArray})

Elementwise multiplication into `dst` of either a scalar or an `NDArray` of the same shape.
Inplace updating.
"""
function mul_to!(dst :: NDArray, arg :: Union{Real, NDArray})
  @assert dst.writable
  if isa(arg, Real)
    _mul_scalar(dst, convert(eltype(dst), arg), dst)
  else
    _mul(dst, arg, dst)
  end
  return dst
end

import Base: .*, *

"""
    .*(arg0, arg1)

Elementwise multiplication of `arg0` and `arg`, could be either scalar or `NDArray`.
"""
function .*(arg0 :: NDArray, arg :: Union{Real, NDArray})
  ret = copy(arg0, context(arg0))
  mul_to!(ret, arg)
end
function .*(arg0 :: Real, arg :: NDArray)
  .*(arg, arg0)
end

"""
    *(arg0, arg1)

Currently only multiplication a scalar with an `NDArray` is implemented. Matrix multiplication
is to be added soon.
"""
function *(arg0 :: NDArray, arg :: Real)
  ret = copy(arg0, context(arg0))
  mul_to!(ret, arg)
end
function *(arg0 :: Real, arg :: NDArray)
  *(arg, arg0)
end

"""
    div_from!(dst :: NDArray, arg :: Union{Real, NDArray})

Elementwise divide a scalar or an `NDArray` of the same shape from `dst`. Inplace updating.
"""
function div_from!(dst :: NDArray, arg :: Union{Real, NDArray})
  @assert dst.writable
  if isa(arg, Real)
    _div_scalar(dst, convert(eltype(dst), arg), dst)
  else
    _div(dst, arg, dst)
  end
end

import Base: ./, /
"""
    ./(arg0 :: NDArray, arg :: Union{Real, NDArray})

Elementwise dividing an `NDArray` by a scalar or another `NDArray` of the same shape.
"""
function ./(arg0 :: NDArray, arg :: Union{Real, NDArray})
  ret = copy(arg0, context(arg0))
  div_from!(ret, arg)
end

"""
    /(arg0 :: NDArray, arg :: Real)

Divide an `NDArray` by a scalar. Matrix division (solving linear systems) is not implemented yet.
"""
function /(arg0 :: NDArray, arg :: Real)
  ./(arg0, arg)
end

"""
Manipulating as Julia Arrays
----------------------------

    @nd_as_jl(captures..., statement)

A convenient macro that allows to operate `NDArray` as Julia Arrays. For example,

```julia
  x = mx.zeros(3,4)
  y = mx.ones(3,4)
  z = mx.zeros((3,4), mx.gpu())

  @mx.nd_as_jl ro=(x,y) rw=z begin
    # now x, y, z are just ordinary Julia Arrays
    z[:,1] = y[:,2]
    z[:,2] = 5
  end
```

Under the hood, the macro convert all the declared captures from `NDArray` into Julia
Arrays, by using `try_get_shared`. And automatically commit the modifications back into
the `NDArray` that is declared as `rw`. This is useful for fast prototyping and when
implement non-critical computations, such as `AbstractEvalMetric`.

!!! note
* Multiple `rw` and / or `ro` capture declaration could be made.
* The macro does **not** check to make sure that `ro` captures are not modified. If the
  original `NDArray` lives in CPU memory, then it is very likely the corresponding
  Julia Array shares data with the `NDArray`, so modifying the Julia Array will also
  modify the underlying `NDArray`.
* More importantly, since the `NDArray` is
  asynchronized, we will wait for *writing* for `rw` variables but wait only for *reading*
  in `ro` variables. If we write into those `ro` variables, **and** if the memory is
  shared, racing condition might happen, and the behavior is undefined.
* When an `NDArray` is declared to be captured as `rw`, its contents is always sync
  back in the end.
* The execution results of the expanded macro is always `nothing`.
* The statements are wrapped in a `let`, thus locally introduced new variables will not be
  available after the statements. So you will need to declare the variables before calling the
  macro if needed.
"""
macro nd_as_jl(m_args...)
  @assert(length(m_args) > 0)
  stmts = m_args[end]
  @assert(isa(stmts, Expr) && stmts.head == :block,
          "The last argument should be a statement block (begin-end); but get $stmts")
  stmts = esc(stmts)

  dclrs  = m_args[1:end-1]
  nd_ro  = []
  nd_rw  = []
  nd_all = []
  for declr in dclrs
    @assert(isa(declr, Expr) && declr.head == :(=) && length(declr.args)==2 && declr.args[1] ∈ (:ro,:rw),
            "Invalid declaration, should be rw=(x,y) or ro=z; but get $declr")

    declr_vars = declr.args[2]
    if isa(declr_vars, Symbol)
      declr_vars = (declr_vars,)
    elseif isa(declr_vars, Expr)
      @assert(declr_vars.head ∈ (:tuple, :vect),
              "Capture declaration should be a variable or a tuple of variables; but got $declr_vars")
      declr_vars = declr_vars.args
    else
      @assert(false, "Capture declaration should be a variable or a tuple of variables; but got $declr_vars")
    end
    for declr_var in declr_vars
      @assert(isa(declr_var, Symbol),
              "Captured ndarrays in ro/rw declaration should be variables, but get $(declr_var)")
    end
    append!(nd_all, [declr_vars...])
    if declr.args[1] == :ro
      append!(nd_ro, [declr_vars...])
    else
      append!(nd_rw, [declr_vars...])
    end
  end

  nd_ro    = map(esc, nd_ro)
  nd_rw    = map(esc, nd_rw)
  nd_all   = map(esc, nd_all)
  rw_origs = [gensym() for _ in nd_rw]

  save_statements  = Expr(:block, [:($v_orig = $v) for (v_orig, v) in zip(rw_origs, nd_rw)]...)
  wait_statements  = Expr(:block, [:(_wait_to_read($v)) for v in nd_ro]...,
                                  [:(_wait_to_write($v)) for v in nd_rw]...)
  clear_statements = Expr(:block, [:($v_orig = nothing) for v_orig in rw_origs]...)
  let_assignments  = [:($v = try_get_shared($v)) for v in nd_all]
  sync_statements  = map(rw_origs, nd_rw) do v_orig, v
    quote
      if !is_shared($v, $v_orig)
        # copy data back if not or no longer sharing data
        copy!($v_orig, $v)
      end
    end
  end
  sync_statements  = Expr(:block, sync_statements...)

  let_statement = Expr(:let, quote
    $stmts
    $sync_statements
  end, let_assignments...)
  m_body = quote
    $wait_statements
    $save_statements
    $let_statement
    $clear_statements
    nothing # the final results is always nothing
  end

  m_body
end

# NOTE: internal use only. Accessing pointers on a different device (e.g. accessing GPU
# pointers from CPU) leads to undefined behavior.
import Base.pointer
function pointer(arr :: NDArray)
  pdata = Ref{Ptr{MX_float}}(0)
  @mxcall(:MXNDArrayGetData, (MX_handle, Ref{Ptr{MX_float}}), arr, pdata)
  return convert(Ptr{eltype(arr)}, pdata[])
end
function _wait_to_read(arr :: NDArray)
  @mxcall(:MXNDArrayWaitToRead, (MX_handle,), arr)
end
function _wait_to_write(arr :: NDArray)
  @mxcall(:MXNDArrayWaitToWrite, (MX_handle,), arr)
end

"""
    try_get_shared(arr)

Try to create a Julia array by sharing the data with the underlying `NDArray`.

# Arguments:
* `arr::NDArray`: the array to be shared.

!!! note
    The returned array does not guarantee to share data with the underlying `NDArray`.
    In particular, data sharing is possible only when the `NDArray` lives on CPU.
"""
function try_get_shared(arr :: NDArray)
  if context(arr).device_type == CPU
    # try to do data sharing
    return unsafe_wrap(Array, pointer(arr), size(arr))
  else
    # impossible to share, just copying
    return copy(arr)
  end
end

"""
    is_shared(j_arr, arr)

Test whether `j_arr` is sharing data with `arr`.

# Arguments:
* Array j_arr: the Julia Array.
* NDArray arr: the `NDArray`.
"""
function is_shared(j_arr :: Array, arr :: NDArray)
  false
end
function is_shared{T<:DType}(j_arr :: Array{T}, arr :: NDArray)
  if length(j_arr) != length(arr)
    return false
  end
  if context(arr).device_type != CPU
    return false
  end
  return pointer(j_arr) == pointer(arr)
end

"""
    load(filename, ::Type{NDArray})

Load NDArrays from binary file.

# Arguments:
* `filename::String`: the path of the file to load. It could be S3 or HDFS address.

Returns either `Dict{Symbol, NDArray}` or `Vector{NDArray}`.

`filename` can point to `s3` or `hdfs` resources if the `libmxnet` is built with the
corresponding components enabled. Examples:
* `s3://my-bucket/path/my-s3-ndarray`
* `hdfs://my-bucket/path/my-hdfs-ndarray`
* `/path-to/my-local-ndarray`
"""
function load(filename::AbstractString, ::Type{NDArray})
  out_size      = Ref{MX_uint}(0)
  out_hdrs      = Ref{Ptr{MX_handle}}(0)
  out_name_size = Ref{MX_uint}(0)
  out_names     = Ref{char_pp}(0)
  @mxcall(:MXNDArrayLoad, (char_p, Ref{MX_uint}, Ref{Ptr{MX_handle}}, Ref{MX_uint}, Ref{char_pp}),
          filename, out_size, out_hdrs, out_name_size, out_names)
  out_name_size = out_name_size[]
  out_size      = out_size[]
  if out_name_size == 0
    return [NDArray(MX_NDArrayHandle(hdr)) for hdr in unsafe_wrap(Array, out_hdrs[], out_size)]
  else
    @assert out_size == out_name_size
    return Dict([(Symbol(unsafe_wrap(String, k)), NDArray(MX_NDArrayHandle(hdr))) for (k,hdr) in
                 zip(unsafe_wrap(Array, out_names[], out_size), unsafe_wrap(Array, out_hdrs[], out_size))])
  end
end

"""
    save(filename :: AbstractString, data)

Save NDarrays to binary file. Filename could be S3 or HDFS address, if `libmxnet` is built
with corresponding support (see `load`).

* `filename::String`: path to the binary file to write to.
* `data`: data to save to file. Data can be a`NDArray`, a `Vector{NDArray}`, or a `Dict{Base.Symbol, NDArray}`.
"""
function save(filename::String, data::NDArray)
  save(filename, [data])
end
function save(filename::String, data::Vector{NDArray})
  @mxcall(:MXNDArraySave, (char_p, MX_uint, Ptr{MX_handle}, char_pp),
          filename, length(data), MX_handle[data...], char_pp(0))
end
function save(filename::String, data::Dict{Base.Symbol,NDArray})
  names  = [k for k in keys(data)]
  arrays = MX_handle[data[k] for k in names]
  names  = String[string(k) for k in names]

  @mxcall(:MXNDArraySave, (char_p, MX_uint, Ptr{MX_handle}, char_pp),
          filename, length(names), arrays, names)
end

################################################################################
# NDArray functions dynamically imported from libmxnet
################################################################################
function _invoke_mxfunction(func_handle::MX_handle, use_vars, scalars, mut_vars; kwargs...)
  names = String[string(entry[1]) for entry in kwargs]
  args = String[string(entry[2]) for entry in kwargs]
  @mxcall(:MXFuncInvokeEx,
          (MX_handle, Ptr{MX_handle}, Ptr{MX_float}, Ptr{MX_handle}, Cint, char_pp, char_pp),
          func_handle, use_vars, scalars, mut_vars, length(names), names, args)
end

@enum(LIBMX_FUNC_TYPE_MASK,
  NDARRAY_ARG_BEFORE_SCALAR = 1,
  ACCEPT_EMPTY_MUTATE_TARGET = (1 << 2)
)

# Import corresponding math functions from base so the automatically defined libmxnet
# functions can overload them
import Base: sqrt

"""
The libxmnet APIs are automatically imported from `libmxnet.so`. The functions listed
here operate on `NDArray` objects. The arguments to the functions are typically ordered
as

```julia
  func_name(arg_in1, arg_in2, ..., scalar1, scalar2, ..., arg_out1, arg_out2, ...)
```

unless `NDARRAY_ARG_BEFORE_SCALAR` is not set. In this case, the scalars are put before the input arguments:

```julia
  func_name(scalar1, scalar2, ..., arg_in1, arg_in2, ..., arg_out1, arg_out2, ...)
```

If `ACCEPT_EMPTY_MUTATE_TARGET` is set. An overloaded function without the output arguments will also be defined:

```julia
  func_name(arg_in1, arg_in2, ..., scalar1, scalar2, ...)
```

Upon calling, the output arguments will be automatically initialized with empty NDArrays.

Those functions always return the output arguments. If there is only one output (the typical situation), that
object (`NDArray`) is returned. Otherwise, a tuple containing all the outputs will be returned.
"""

function _get_ndarray_functions()
  n = Ref{MX_uint}(0)
  handles = Ref{Ptr{MX_handle}}(0)

  @mxcall(:MXListFunctions, (Ref{MX_uint}, Ref{Ptr{MX_handle}}), n, handles)

  funcs = unsafe_wrap(Array, handles[], n[])
  return funcs
end

const _function_cache = Dict{Symbol, MX_handle}()
function _get_function(name :: Symbol)
  if !haskey(_function_cache, name)
    handle = Ref{MX_handle}(0)

    @mxcall(:MXGetFunction, (Cstring, Ref{MX_handle}), name, handle)
    _function_cache[name] = handle[]
    return handle[]
  else
    return _function_cache[name]
  end
end

function _get_function_description(handle :: MX_handle)
  # get function information (human readable)
  ref_name = Ref{char_p}(0)
  ref_desc = Ref{char_p}(0)
  ref_narg = Ref{MX_uint}(0)

  ref_arg_names = Ref{char_pp}(0)
  ref_arg_types = Ref{char_pp}(0)
  ref_arg_descs = Ref{char_pp}(0)

  ref_ret_type  = Ref{char_p}(0)

  @mxcall(:MXFuncGetInfo,
         (MX_handle, Ref{char_p}, Ref{char_p}, Ref{MX_uint}, Ref{char_pp},
          Ref{char_pp}, Ref{char_pp}, Ref{char_p}),
          handle, ref_name, ref_desc, ref_narg, ref_arg_names,
          ref_arg_types, ref_arg_descs, ref_ret_type)

  name = Symbol(unsafe_wrap(String, ref_name[]))
  signature = _format_signature(Int(ref_narg[]), ref_arg_names)
  desc = "    " * string(name) * "(" * signature * ")\n\n"
  desc *= unsafe_wrap(String, ref_desc[]) * "\n\n"
  desc *= "# Arguments\n"
  desc *= _format_docstring(Int(ref_narg[]), ref_arg_names, ref_arg_types, ref_arg_descs)
  return name, desc
end

function _get_function_expressions(handle :: MX_handle, name)
  # get function specification
  ref_n_use_vars = Ref{MX_uint}(0)
  ref_n_scalars  = Ref{MX_uint}(0)
  ref_n_mut_vars = Ref{MX_uint}(0)
  ref_type_mask  = Ref{Cint}(0)
  @mxcall(:MXFuncDescribe,
         (MX_handle, Ref{MX_uint}, Ref{MX_uint}, Ref{MX_uint}, Ref{Cint}),
          handle, ref_n_use_vars, ref_n_scalars, ref_n_mut_vars, ref_type_mask)

  n_used_vars   = ref_n_use_vars[]
  n_scalars     = ref_n_scalars[]
  n_mutate_vars = ref_n_mut_vars[]
  type_mask     = ref_type_mask[]
  accept_empty_mutate = (type_mask & convert(Cint,ACCEPT_EMPTY_MUTATE_TARGET)) != 0
  arg_before_scalar   = (type_mask & convert(Cint,NDARRAY_ARG_BEFORE_SCALAR)) != 0

  # general ndarray function
  if arg_before_scalar
    args = vcat([Expr(:(::), Symbol("in$i"), NDArray) for i=1:n_used_vars],
                [Expr(:(::), Symbol("sca$i"), Real) for i=1:n_scalars],
                [Expr(:(::), Symbol("out$i"), NDArray) for i=1:n_mutate_vars])
  else
    args = vcat([Expr(:(::), Symbol("sca$i"), Real) for i=1:n_scalars],
                [Expr(:(::), Symbol("in$i"), NDArray) for i=1:n_used_vars],
                [Expr(:(::), Symbol("out$i"), NDArray) for i=1:n_mutate_vars])
  end

  _use_vars = Expr(:ref, :MX_handle, [Symbol("in$i") for i=1:n_used_vars]...)
  _scalars  = Expr(:ref, :MX_float, [Symbol("sca$i") for i=1:n_scalars]...)
  _mut_vars = Expr(:ref, :MX_handle, [Symbol("out$i") for i=1:n_mutate_vars]...)

  # XXX: hacky way of solving the problem that the arguments of `dot` should be swapped
  # See https://github.com/dmlc/MXNet.jl/issues/55
  if name == :dot
    _use_vars.args[2:end] = flipdim(_use_vars.args[2:end], 1)
  end

  # XXX: hacky way of solving the semantic difference of the axes parameter in Julia
  # and in libmxnet.
  # See https://github.com/dmlc/MXNet.jl/pull/123
  if name == :transpose
    transform = quote
      kwargs = Any[key != :axes ? (key, arg) : (key, reverse(map(i->length(arg)-i, arg))) for (key, arg) in kwargs]
    end
  else
    transform = :()
  end

  stmt_call = quote
    local handle = _get_function($(QuoteNode(name)))
    _invoke_mxfunction(handle, $_use_vars, $_scalars, $_mut_vars; kwargs...)
  end
  if n_mutate_vars == 1
    stmt_ret = :(return out1)
  else
    stmt_ret = Expr(:return, Expr(:tuple, [Symbol("out$i") for i=1:n_mutate_vars]...))
  end

  func_def = quote
    function $name($(args...); kwargs...)
      $transform
      $stmt_call
      $stmt_ret
    end
  end

  if accept_empty_mutate
    args0      = args[1:n_used_vars+n_scalars]
    _mut_vars0 = [:(NDArray(_ndarray_alloc())) for i=1:n_mutate_vars]

    func_def0 = quote
      function $name($(args0...); kwargs...)
        $name($(args0...), $(_mut_vars0...); kwargs...)
      end
    end
    return func_def, func_def0
  else
    return func_def, :()
  end
end

macro _import_ndarray_functions()
  funcs = _get_ndarray_functions()
  func_exprs = Expr[]

  for i = 1:length(funcs)
    handle = funcs[i]

    name, desc = _get_function_description(handle)
    func_def, func_def0 = _get_function_expressions(handle, name)

    expr = quote
      $(isdefined(Base, name) ? :(import Base.$name) : :())
      @doc $desc ->
      $func_def
      $func_def0
    end

    push!(func_exprs, expr)
  end

  esc(quote
    $(func_exprs...)
  end)
end

@_import_ndarray_functions()